diff --git a/.env.example b/.env.example index 30535834..765ea3f6 100644 --- a/.env.example +++ b/.env.example @@ -70,6 +70,12 @@ NEARAI_AUTH_URL=https://private.near.ai # LLM_BASE_URL=https://api.fireworks.ai/inference/v1 # LLM_API_KEY=fw_... +# === MiniMax === +# LLM_BACKEND=minimax +# MINIMAX_API_KEY=... +# MINIMAX_MODEL=MiniMax-M2.5 +# MINIMAX_BASE_URL=https://api.minimax.io/v1 # default (global); use https://api.minimaxi.com/v1 for China + # === Anthropic Direct === # LLM_BACKEND=anthropic # ANTHROPIC_MODEL=claude-sonnet-4-6 diff --git a/.github/scripts/pr-body-utils.sh b/.github/scripts/pr-body-utils.sh new file mode 100644 index 00000000..f41f769f --- /dev/null +++ b/.github/scripts/pr-body-utils.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash + +load_commit_summary() { + local range="$1" + local max_commits="${2:-50}" + local commit_list overflow + + commit_list="$(git log --oneline --no-merges --reverse "${range}" 2>/dev/null || echo "")" + if [ -n "${commit_list}" ]; then + COMMIT_COUNT="$(printf '%s\n' "${commit_list}" | wc -l | tr -d ' ')" + if [ "${COMMIT_COUNT}" -gt "${max_commits}" ]; then + COMMIT_MD="$(printf '%s\n' "${commit_list}" | head -n "${max_commits}" | sed 's/^/- /')" + overflow=$((COMMIT_COUNT - max_commits)) + COMMIT_MD+=$'\n'"- ... and ${overflow} more (see compare view)" + else + COMMIT_MD="$(printf '%s\n' "${commit_list}" | sed 's/^/- /')" + fi + else + COMMIT_COUNT=0 + COMMIT_MD="- (no non-merge commits in range)" + fi +} + +replace_marked_section() { + local body_file="$1" + local section_file="$2" + local section_start="$3" + local section_end="$4" + local output_file="$5" + + if grep -qF "${section_start}" "${body_file}" && grep -qF "${section_end}" "${body_file}"; then + awk -v start="${section_start}" -v end="${section_end}" -v replacement_file="${section_file}" ' + BEGIN { + while ((getline line < replacement_file) > 0) { + replacement = replacement line ORS + } + in_block = 0 + } + $0 == start { + printf "%s", replacement + in_block = 1 + next + } + $0 == end { + in_block = 0 + next + } + !in_block { + print + } + ' "${body_file}" > "${output_file}" + else + cp "${body_file}" "${output_file}" + if [ -s "${output_file}" ]; then + printf '\n\n' >> "${output_file}" + fi + cat "${section_file}" >> "${output_file}" + fi +} diff --git a/.github/scripts/update-release-plz-body.sh b/.github/scripts/update-release-plz-body.sh new file mode 100644 index 00000000..3a7eef20 --- /dev/null +++ b/.github/scripts/update-release-plz-body.sh @@ -0,0 +1,101 @@ +#!/usr/bin/env bash +set -euo pipefail + +: "${PR_NUMBER:?PR_NUMBER is required}" +: "${REPO:?REPO is required}" + +MAIN_BRANCH="${MAIN_BRANCH:-main}" +DRY_RUN="${DRY_RUN:-false}" +SECTION_START="" +SECTION_END="" +TMP_DIR="$(mktemp -d)" +trap 'rm -rf "${TMP_DIR}"' EXIT + +# shellcheck source=.github/scripts/pr-body-utils.sh +source "$(dirname "$0")/pr-body-utils.sh" + +gh pr view "${PR_NUMBER}" --repo "${REPO}" --json body > "${TMP_DIR}/pr.json" +jq -r '.body // ""' < "${TMP_DIR}/pr.json" > "${TMP_DIR}/body.md" + +git fetch origin "${MAIN_BRANCH}" +git fetch origin "+refs/tags/v*:refs/tags/v*" + +LAST_TAG="$(git describe --tags --match 'v*' --abbrev=0 "origin/${MAIN_BRANCH}" 2>/dev/null || true)" +if [ -n "${LAST_TAG}" ]; then + RANGE="${LAST_TAG}..origin/${MAIN_BRANCH}" + HEADER="## Staging promotion batches since ${LAST_TAG}" + EMPTY_MESSAGE="_No structured staging promotion merges found since ${LAST_TAG}._" +else + RANGE="origin/${MAIN_BRANCH}" + HEADER="## Staging promotion batches on ${MAIN_BRANCH}" + EMPTY_MESSAGE="_No structured staging promotion merges found on ${MAIN_BRANCH}._" +fi + +{ + echo "${SECTION_START}" + echo "${HEADER}" + echo +} > "${TMP_DIR}/section.md" + +FOUND_SUMMARY=false +while IFS= read -r sha; do + [ -n "${sha}" ] || continue + BODY="$(git show -s --format=%b "${sha}")" + if ! printf '%s\n' "${BODY}" | grep -q '^staging-promotion-summary-v1$'; then + continue + fi + + FOUND_SUMMARY=true + SUBJECT="$(git show -s --format=%s "${sha}")" + PR_REF="$(printf '%s\n' "${BODY}" | sed -n 's/^promotion-pr: //p' | head -n 1)" + COMMIT_COUNT="$(printf '%s\n' "${BODY}" | sed -n 's/^current-commit-count: //p' | head -n 1)" + CURRENT_RANGE="$(printf '%s\n' "${BODY}" | sed -n 's/^current-range: //p' | head -n 1)" + COMMIT_BLOCK="$(printf '%s\n' "${BODY}" | awk 'capture { print } /^Current commits in this promotion \([0-9]+\):$/ { capture = 1 }')" + + { + echo "### ${SUBJECT}" + echo + if [ -n "${PR_REF}" ]; then + echo "**Promotion PR:** ${PR_REF}" + fi + if [ -n "${COMMIT_COUNT}" ]; then + echo "**Commit count:** ${COMMIT_COUNT}" + fi + if [ -n "${CURRENT_RANGE}" ]; then + echo "**Range:** \`${CURRENT_RANGE}\`" + fi + echo + if [ -n "${COMMIT_BLOCK}" ]; then + echo "${COMMIT_BLOCK}" + else + echo "- (no commit summary found)" + fi + echo + } >> "${TMP_DIR}/section.md" +done < <(git log --merges --reverse --format='%H' "${RANGE}") + +if [ "${FOUND_SUMMARY}" = false ]; then + { + echo "${EMPTY_MESSAGE}" + echo + } >> "${TMP_DIR}/section.md" +fi + +{ + echo "*Auto-updated from structured staging promotion merge bodies on ${MAIN_BRANCH}.*" + echo "${SECTION_END}" +} >> "${TMP_DIR}/section.md" + +replace_marked_section \ + "${TMP_DIR}/body.md" \ + "${TMP_DIR}/section.md" \ + "${SECTION_START}" \ + "${SECTION_END}" \ + "${TMP_DIR}/new-body.md" + +if [ "${DRY_RUN}" = "true" ]; then + echo "Dry run enabled. Computed PR body for #${PR_NUMBER}:" + cat "${TMP_DIR}/new-body.md" +else + gh pr edit "${PR_NUMBER}" --repo "${REPO}" --body-file "${TMP_DIR}/new-body.md" +fi diff --git a/.github/scripts/update-staging-promotion-body.sh b/.github/scripts/update-staging-promotion-body.sh new file mode 100644 index 00000000..9686b58c --- /dev/null +++ b/.github/scripts/update-staging-promotion-body.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +set -euo pipefail + +: "${PR_NUMBER:?PR_NUMBER is required}" +: "${REPO:?REPO is required}" + +MAX_COMMITS="${MAX_COMMITS:-50}" +DRY_RUN="${DRY_RUN:-false}" +SECTION_START="" +SECTION_END="" +TMP_DIR="$(mktemp -d)" +trap 'rm -rf "${TMP_DIR}"' EXIT + +# shellcheck source=.github/scripts/pr-body-utils.sh +source "$(dirname "$0")/pr-body-utils.sh" + +gh pr view "${PR_NUMBER}" --repo "${REPO}" --json body,baseRefName,headRefName > "${TMP_DIR}/pr.json" +jq -r '.body // ""' < "${TMP_DIR}/pr.json" > "${TMP_DIR}/body.md" +BASE="$(jq -r '.baseRefName' < "${TMP_DIR}/pr.json")" +HEAD="$(jq -r '.headRefName' < "${TMP_DIR}/pr.json")" +RANGE="origin/${BASE}..origin/${HEAD}" + +git fetch origin "${BASE}" "${HEAD}" + +load_commit_summary "${RANGE}" "${MAX_COMMITS}" + +{ + echo "${SECTION_START}" + echo "### Current commits in this promotion (${COMMIT_COUNT})" + echo + echo "**Current base:** \`${BASE}\`" + echo "**Current head:** \`${HEAD}\`" + echo "**Current range:** \`${RANGE}\`" + echo + echo "${COMMIT_MD}" + echo + echo "*Auto-updated by staging promotion metadata workflow*" + echo "${SECTION_END}" +} > "${TMP_DIR}/section.md" + +replace_marked_section \ + "${TMP_DIR}/body.md" \ + "${TMP_DIR}/section.md" \ + "${SECTION_START}" \ + "${SECTION_END}" \ + "${TMP_DIR}/new-body.md" + +if [ "${DRY_RUN}" = "true" ]; then + echo "Dry run enabled. Computed PR body for #${PR_NUMBER}:" + cat "${TMP_DIR}/new-body.md" +else + gh pr edit "${PR_NUMBER}" --repo "${REPO}" --body-file "${TMP_DIR}/new-body.md" +fi diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index 01352005..fef89bae 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -52,7 +52,7 @@ jobs: - group: features files: "tests/e2e/scenarios/test_skills.py tests/e2e/scenarios/test_tool_approval.py" - group: extensions - files: "tests/e2e/scenarios/test_extensions.py" + files: "tests/e2e/scenarios/test_extensions.py tests/e2e/scenarios/test_extension_oauth.py tests/e2e/scenarios/test_wasm_lifecycle.py tests/e2e/scenarios/test_tool_execution.py tests/e2e/scenarios/test_pairing.py" steps: - uses: actions/checkout@v6 diff --git a/.github/workflows/release-plz-batch-summary.yml b/.github/workflows/release-plz-batch-summary.yml new file mode 100644 index 00000000..0e106736 --- /dev/null +++ b/.github/workflows/release-plz-batch-summary.yml @@ -0,0 +1,44 @@ +name: Release-plz Batch Summary + +on: + workflow_dispatch: + inputs: + pr_number: + description: "release-plz PR number to refresh" + required: true + type: string + dry_run: + description: "Compute the body update without editing the PR" + required: false + type: boolean + default: true + pull_request_target: + types: [opened, synchronize, reopened] + +permissions: + contents: read + pull-requests: write + +jobs: + update-release-pr: + if: > + (github.event_name == 'pull_request_target' && + github.event.pull_request.head.repo.full_name == github.repository && + startsWith(github.event.pull_request.head.ref, 'release-plz-')) || + github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + steps: + - name: Checkout base branch + uses: actions/checkout@v6 + with: + ref: ${{ github.event_name == 'workflow_dispatch' && 'main' || github.event.pull_request.base.ref }} + fetch-depth: 0 + fetch-tags: true + + - name: Update release-plz PR body with staging batch summary + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + PR_NUMBER: ${{ github.event_name == 'workflow_dispatch' && inputs.pr_number || github.event.pull_request.number }} + REPO: ${{ github.repository }} + DRY_RUN: ${{ github.event_name == 'workflow_dispatch' && inputs.dry_run || 'false' }} + run: bash .github/scripts/update-release-plz-body.sh diff --git a/.github/workflows/release-plz.yml b/.github/workflows/release-plz.yml index 142b2b20..d1be9004 100644 --- a/.github/workflows/release-plz.yml +++ b/.github/workflows/release-plz.yml @@ -58,10 +58,16 @@ jobs: - *checkout - *install-rust - uses: Swatinem/rust-cache@v2 + - name: Generate GitHub token + uses: actions/create-github-app-token@v2 + id: generate-token + with: + app-id: ${{ secrets.GH_RELEASES_MANAGER_APP_ID }} + private-key: ${{ secrets.GH_RELEASES_MANAGER_APP_PRIVATE_KEY }} - name: Run release-plz uses: release-plz/action@v0.5 with: command: release-pr env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_TOKEN: ${{ steps.generate-token.outputs.token }} CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }} diff --git a/.github/workflows/staging-ci.yml b/.github/workflows/staging-ci.yml index c8a39c80..2df7bf6f 100644 --- a/.github/workflows/staging-ci.yml +++ b/.github/workflows/staging-ci.yml @@ -25,9 +25,35 @@ concurrency: cancel-in-progress: false # Let running suites finish jobs: + # ── Resolve promotion base branch ─────────────────────────────── + resolve-promotion-base: + name: Resolve promotion base + runs-on: ubuntu-latest + outputs: + promotion_base: ${{ steps.resolve.outputs.promotion_base }} + steps: + - name: Resolve promotion base + id: resolve + env: + GH_TOKEN: ${{ github.token }} + FALLBACK_BRANCH: main + REPO: ${{ github.repository }} + run: | + LATEST=$(gh pr list --repo "${REPO}" --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 "promotion_base=${LATEST}" >> "$GITHUB_OUTPUT" + echo "Using open promotion branch as base: ${LATEST}" + else + echo "promotion_base=${FALLBACK_BRANCH}" >> "$GITHUB_OUTPUT" + echo "No open promotion branch found. Using ${FALLBACK_BRANCH}." + fi + # ── Check for new commits ────────────────────────────────────── check-changes: name: Check for new commits + needs: resolve-promotion-base runs-on: ubuntu-latest outputs: has_changes: ${{ steps.check.outputs.has_changes }} @@ -44,6 +70,7 @@ jobs: id: check env: FORCE_RUN: ${{ inputs.force }} + PROMOTION_BASE: ${{ needs.resolve-promotion-base.outputs.promotion_base }} run: | CURRENT_HEAD=$(git rev-parse HEAD) echo "current_head=${CURRENT_HEAD}" >> "$GITHUB_OUTPUT" @@ -65,9 +92,9 @@ jobs: 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}" + git fetch origin "${PROMOTION_BASE}" + MERGE_BASE=$(git merge-base "origin/${PROMOTION_BASE}" HEAD) + echo "First run -- reviewing from merge-base ${MERGE_BASE} against ${PROMOTION_BASE}" DIFF_RANGE="${MERGE_BASE}..${CURRENT_HEAD}" fi fi @@ -101,7 +128,7 @@ jobs: # ── Create promotion PR (triggers claude-review.yml on the PR) ── create-promotion-pr: name: Create Promotion PR - needs: check-changes + needs: [resolve-promotion-base, check-changes] if: needs.check-changes.outputs.has_changes == 'true' runs-on: ubuntu-latest outputs: @@ -129,18 +156,19 @@ jobs: echo "token=${{ github.token }}" >> "$GITHUB_OUTPUT" fi - - name: Check if staging is ahead of main + - name: Check if staging is ahead of target branch id: ahead-check env: GH_TOKEN: ${{ steps.token.outputs.token }} + PROMOTION_BASE: ${{ needs.resolve-promotion-base.outputs.promotion_base }} run: | - git fetch origin main - AHEAD=$(git rev-list --count origin/main..origin/staging) + git fetch origin "${PROMOTION_BASE}" + AHEAD=$(git rev-list --count "origin/${PROMOTION_BASE}..origin/staging") echo "commits_ahead=${AHEAD}" >> "$GITHUB_OUTPUT" if [ "$AHEAD" -eq 0 ]; then - echo "Staging is not ahead of main. Nothing to promote." + echo "Staging is not ahead of ${PROMOTION_BASE}. Nothing to promote." else - echo "Staging is ${AHEAD} commits ahead of main." + echo "Staging is ${AHEAD} commits ahead of ${PROMOTION_BASE}." fi - name: Create promotion branch @@ -154,53 +182,53 @@ jobs: 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: | + source .github/scripts/pr-body-utils.sh 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 }}" + BASE="${{ needs.resolve-promotion-base.outputs.promotion_base }}" + + MAX_COMMITS=50 + load_commit_summary "${RANGE}" "${MAX_COMMITS}" + + # Build PR body via concatenation to avoid heredoc shell expansion + # (commit messages in COMMIT_MD may contain $, backticks, or backslashes) + PR_BODY="## Auto-promotion from staging CI" + PR_BODY+=$'\n\n'"**Batch range:** \`${RANGE}\`" + PR_BODY+=$'\n'"**Promotion branch:** \`${BRANCH}\`" + PR_BODY+=$'\n'"**Base:** \`${BASE}\`" + PR_BODY+=$'\n'"**Triggered by:** Staging CI batch at ${TIMESTAMP}" + PR_BODY+=$'\n\n'"### Commits in this batch (${COMMIT_COUNT}):" + PR_BODY+=$'\n'"${COMMIT_MD}" + PR_BODY+=$'\n\n'"" + PR_BODY+=$'\n'"### Current commits in this promotion (${COMMIT_COUNT})" + PR_BODY+=$'\n' + PR_BODY+=$'\n'"**Current base:** \`${BASE}\`" + PR_BODY+=$'\n'"**Current head:** \`${BRANCH}\`" + PR_BODY+=$'\n'"**Current range:** \`origin/${BASE}..origin/${BRANCH}\`" + PR_BODY+=$'\n' + PR_BODY+=$'\n'"${COMMIT_MD}" + PR_BODY+=$'\n' + PR_BODY+=$'\n'"*Auto-updated by staging promotion metadata workflow*" + PR_BODY+=$'\n'"" + PR_BODY+=$'\n\n'"Waiting for gates:" + PR_BODY+=$'\n'"- Tests: pending" + PR_BODY+=$'\n'"- E2E: pending" + PR_BODY+=$'\n'"- Claude Code review: pending (will post comments on this PR)" + PR_BODY+=$'\n\n'"---" + PR_BODY+=$'\n'"*Auto-created by staging-ci workflow*" 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*" \ + --title "chore: promote staging to ${BASE} (${TIMESTAMP})" \ + --body "$PR_BODY" \ --label "staging-promotion") PR_NUM=$(echo "$PR_URL" | grep -oE '[0-9]+$') @@ -225,7 +253,8 @@ jobs: - uses: actions/checkout@v6 with: ref: staging - fetch-depth: 1 + # Need full history to recompute the final promoted range before merge. + fetch-depth: 0 - name: Generate GitHub App token id: app-token @@ -324,8 +353,10 @@ jobs: # 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/') + SEVERITY="${TAG#\[}" + SEVERITY="${SEVERITY%%:*}" + CONFIDENCE="${TAG##*:}" + CONFIDENCE="${CONFIDENCE%\]}" DESC=$(echo "$line" | sed "s/\[${SEVERITY}:${CONFIDENCE}\] *//" | head -1) echo "Found: [${SEVERITY}:${CONFIDENCE}] ${DESC}" @@ -417,11 +448,29 @@ jobs: GH_TOKEN: ${{ steps.token.outputs.token }} PR_NUMBER: ${{ needs.create-promotion-pr.outputs.pr_number }} run: | + source .github/scripts/pr-body-utils.sh if [ -n "$PR_NUMBER" ]; then BASE=$(gh pr view "$PR_NUMBER" --json baseRefName --jq '.baseRefName') if [ "$BASE" = "main" ]; then echo "Merging promotion PR #${PR_NUMBER} (targets main)" - gh pr merge "$PR_NUMBER" --merge + TITLE=$(gh pr view "$PR_NUMBER" --json title --jq '.title') + HEAD_BRANCH=$(gh pr view "$PR_NUMBER" --json headRefName --jq '.headRefName') + git fetch origin "${BASE}" "${HEAD_BRANCH}" + CURRENT_RANGE="origin/${BASE}..origin/${HEAD_BRANCH}" + MAX_COMMITS=50 + load_commit_summary "${CURRENT_RANGE}" "${MAX_COMMITS}" + { + echo "staging-promotion-summary-v1" + echo "promotion-pr: #${PR_NUMBER}" + echo "base: ${BASE}" + echo "head: ${HEAD_BRANCH}" + echo "current-range: ${CURRENT_RANGE}" + echo "current-commit-count: ${COMMIT_COUNT}" + echo "" + echo "Current commits in this promotion (${COMMIT_COUNT}):" + echo "${COMMIT_MD}" + } > /tmp/staging-promotion-merge-body.md + gh pr merge "$PR_NUMBER" --merge --subject "#${PR_NUMBER} $TITLE" --body-file /tmp/staging-promotion-merge-body.md echo "merged=true" >> "$GITHUB_OUTPUT" else echo "PR #${PR_NUMBER} targets '${BASE}' (not main) — leaving open for chain resolution" @@ -461,18 +510,20 @@ jobs: 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 + { + echo "## Staging CI Batch Results" + echo "" + echo "| Check | Result |" + echo "|-------|--------|" + echo "| Tests | ${{ needs.tests.result }} |" + echo "| E2E | ${{ needs.e2e.result }} |" + echo "| Promotion PR | ${{ needs.create-promotion-pr.result }} |" + echo "| Gate | ${{ needs.gate.result }} |" + echo "| Tag Updated | ${{ needs.update-tag.result }} |" + echo "" + echo "Range: ${{ needs.check-changes.outputs.diff_range }}" + PR_NUM="${{ needs.create-promotion-pr.outputs.pr_number }}" + if [ -n "$PR_NUM" ]; then + echo "Promotion PR: #${PR_NUM}" + fi + } >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/workflows/staging-promotion-metadata.yml b/.github/workflows/staging-promotion-metadata.yml new file mode 100644 index 00000000..76b8326b --- /dev/null +++ b/.github/workflows/staging-promotion-metadata.yml @@ -0,0 +1,78 @@ +name: Staging Promotion Metadata + +on: + workflow_dispatch: + inputs: + pr_number: + description: "Staging promotion PR number to refresh" + required: true + type: string + dry_run: + description: "Compute the body update without editing the PR" + required: false + type: boolean + default: true + pull_request_target: + types: [opened, synchronize, reopened] + push: + branches: + - main + +permissions: + contents: read + pull-requests: write + +jobs: + refresh-single-pr: + if: > + (github.event_name == 'pull_request_target' && + github.event.pull_request.head.repo.full_name == github.repository && + startsWith(github.event.pull_request.head.ref, 'staging-promote/')) || + github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + steps: + - name: Checkout workflow source + uses: actions/checkout@v6 + with: + # For chained promotion PRs, the script lives on the trusted PR head, + # not necessarily on the older promotion branch used as the PR base. + ref: ${{ github.event_name == 'workflow_dispatch' && 'main' || github.event.pull_request.head.sha }} + fetch-depth: 0 + fetch-tags: true + + - name: Refresh staging promotion PR body + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + PR_NUMBER: ${{ github.event_name == 'workflow_dispatch' && inputs.pr_number || github.event.pull_request.number }} + REPO: ${{ github.repository }} + DRY_RUN: ${{ github.event_name == 'workflow_dispatch' && inputs.dry_run || 'false' }} + run: bash .github/scripts/update-staging-promotion-body.sh + + refresh-open-prs-after-main-push: + if: github.event_name == 'push' + runs-on: ubuntu-latest + steps: + - name: Checkout main + uses: actions/checkout@v6 + with: + ref: main + fetch-depth: 0 + fetch-tags: true + + - name: Refresh all open staging promotion PR bodies + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + REPO: ${{ github.repository }} + run: | + # ubuntu-latest uses bash 5.x, so mapfile is available here. + mapfile -t prs < <(gh pr list --repo "${REPO}" --label staging-promotion --state open \ + --json number,headRefName \ + --jq '.[] | select(.headRefName | startswith("staging-promote/")) | .number') + if [ "${#prs[@]}" -eq 0 ]; then + echo "No open staging promotion PRs to refresh." + exit 0 + fi + for pr in "${prs[@]}"; do + echo "Refreshing staging promotion PR #${pr}" + PR_NUMBER="${pr}" bash .github/scripts/update-staging-promotion-body.sh + done diff --git a/Dockerfile b/Dockerfile index 0375e509..08a0b721 100644 --- a/Dockerfile +++ b/Dockerfile @@ -19,6 +19,7 @@ WORKDIR /app # Copy manifests first for layer caching COPY Cargo.toml Cargo.lock ./ +COPY crates/ crates/ # Copy source, build script, tests, and supporting directories COPY build.rs build.rs diff --git a/Dockerfile.test b/Dockerfile.test index 202bd04d..6ec502ba 100644 --- a/Dockerfile.test +++ b/Dockerfile.test @@ -20,6 +20,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ WORKDIR /app COPY Cargo.toml Cargo.lock ./ +COPY crates/ crates/ COPY build.rs build.rs COPY src/ src/ COPY tests/ tests/ diff --git a/FEATURE_PARITY.md b/FEATURE_PARITY.md index e53929ca..323a5a38 100644 --- a/FEATURE_PARITY.md +++ b/FEATURE_PARITY.md @@ -170,7 +170,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O | `nodes` | ✅ | ❌ | P3 | Device management, remove/clear flows | | `plugins` | ✅ | ❌ | P3 | Plugin management | | `hooks` | ✅ | ✅ | P2 | Lifecycle hooks | -| `cron` | ✅ | ❌ | P2 | Scheduled jobs (model/thinking fields in edit) | +| `cron` | ✅ | 🚧 | P2 | list/create/edit/enable/disable/delete/history; TODO: `cron run`, model/thinking fields | | `webhooks` | ✅ | ❌ | P3 | Webhook config | | `message send` | ✅ | ❌ | P2 | Send to channels | | `browser` | ✅ | ❌ | P3 | Browser automation | diff --git a/channels-src/discord/Cargo.lock b/channels-src/discord/Cargo.lock index e3a81af1..f25ce551 100644 --- a/channels-src/discord/Cargo.lock +++ b/channels-src/discord/Cargo.lock @@ -20,12 +20,27 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + [[package]] name = "bitflags" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -33,20 +48,134 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] -name = "discord-channel" -version = "0.1.0" +name = "const-oid" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "curve25519-dalek" +version = "4.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" +dependencies = [ + "cfg-if", + "cpufeatures", + "curve25519-dalek-derive", + "digest", + "fiat-crypto", + "rustc_version", + "subtle", + "zeroize", +] + +[[package]] +name = "curve25519-dalek-derive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "der" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" +dependencies = [ + "const-oid", + "zeroize", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "discord-channel" +version = "0.2.0" +dependencies = [ + "ed25519-dalek", + "hex", "serde", "serde_json", "wit-bindgen", ] +[[package]] +name = "ed25519" +version = "2.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" +dependencies = [ + "pkcs8", + "signature", +] + +[[package]] +name = "ed25519-dalek" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" +dependencies = [ + "curve25519-dalek", + "ed25519", + "serde", + "sha2", + "subtle", + "zeroize", +] + [[package]] name = "equivalent" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "fiat-crypto" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "hashbrown" version = "0.14.5" @@ -68,6 +197,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "id-arena" version = "2.3.0" @@ -98,6 +233,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "884e2677b40cc8c339eaefcb701c32ef1fd2493d71118dc0ca4b6a736c93bd67" +[[package]] +name = "libc" +version = "0.2.182" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112" + [[package]] name = "log" version = "0.4.29" @@ -116,6 +257,16 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "pkcs8" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" +dependencies = [ + "der", + "spki", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -144,6 +295,15 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "semver" version = "1.0.27" @@ -193,6 +353,23 @@ dependencies = [ "zmij", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" + [[package]] name = "smallvec" version = "1.15.1" @@ -208,6 +385,22 @@ dependencies = [ "smallvec", ] +[[package]] +name = "spki" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" +dependencies = [ + "base64ct", + "der", +] + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + [[package]] name = "syn" version = "2.0.117" @@ -219,6 +412,12 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "typenum" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" + [[package]] name = "unicode-ident" version = "1.0.24" @@ -394,6 +593,12 @@ dependencies = [ "syn", ] +[[package]] +name = "zeroize" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" + [[package]] name = "zmij" version = "1.0.21" diff --git a/channels-src/discord/Cargo.toml b/channels-src/discord/Cargo.toml index 81e95260..a2892494 100644 --- a/channels-src/discord/Cargo.toml +++ b/channels-src/discord/Cargo.toml @@ -10,6 +10,8 @@ publish = false serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" wit-bindgen = "0.36" +ed25519-dalek = { version = "2", default-features = false, features = ["alloc", "fast", "zeroize"] } +hex = "0.4" [lib] crate-type = ["cdylib"] diff --git a/channels-src/discord/README.md b/channels-src/discord/README.md index 6cb0199f..333e7670 100644 --- a/channels-src/discord/README.md +++ b/channels-src/discord/README.md @@ -21,11 +21,10 @@ WASM channel for Discord integration - handle slash commands and button interact ironclaw secret set discord_bot_token YOUR_BOT_TOKEN ``` - **Note:** The `discord_bot_token` secret is the only value read directly by this - Discord channel WASM component. The `discord_app_id` and `discord_public_key` - secrets are used by the IronClaw host (for example, to verify Discord - interaction signatures and manage slash command registration) and are not - accessed from the WASM module itself. + **Note:** The `discord_bot_token` secret is used for Discord REST API calls. + Interaction signature verification is performed inside the Discord channel + module and uses the channel config field `webhook_secret` (set this to your + Discord app public key hex). ## Discord Configuration @@ -87,6 +86,30 @@ If an internal error occurs (e.g., metadata serialization failure), the tool att Check the host logs for detailed error information. ## Advanced Usage +### Mention Polling + +The Discord channel can also poll configured channels for `@bot` mentions. + +Example channel config: + +```json +{ + "require_signature_verification": true, + "webhook_secret": "YOUR_DISCORD_PUBLIC_KEY_HEX", + "polling_enabled": true, + "poll_interval_ms": 30000, + "mention_channel_ids": ["123456789012345678"], + "owner_id": null, + "dm_policy": "pairing", + "allow_from": [] +} +``` + +### Access Control + +- `owner_id`: when set, only that Discord user can interact with the bot. +- `dm_policy`: `open` allows all DMs; `pairing` requires approval. +- `allow_from`: allowlist entries for DM pairing checks (`*`, user id, or username). ### Embeds @@ -96,8 +119,11 @@ To send embeds, include an `embeds` array in the `metadata_json` field of the ag ### "Invalid Signature" -- Check that `discord_public_key` is set correctly in IronClaw secrets. -- This validation happens on the host before reaching the WASM. +- Check that `webhook_secret` is set to your Discord app public key hex in the + Discord channel config. +- Validation happens inside the Discord WASM channel. +- If `require_signature_verification` is `true` and `webhook_secret` is empty, + the channel returns HTTP `500` with a configuration error. ### "401 Unauthorized" diff --git a/channels-src/discord/discord.capabilities.json b/channels-src/discord/discord.capabilities.json index fd55c685..9ff7a890 100644 --- a/channels-src/discord/discord.capabilities.json +++ b/channels-src/discord/discord.capabilities.json @@ -3,7 +3,7 @@ "wit_version": "0.3.0", "type": "channel", "name": "discord", - "description": "Discord Gateway/Webhook channel for handling slash commands, buttons, and messages", + "description": "Discord webhook channel for slash commands, components, and optional mention polling", "setup": { "required_secrets": [ { @@ -41,7 +41,7 @@ }, "channel": { "allowed_paths": ["/webhook/discord"], - "allow_polling": false, + "allow_polling": true, "callback_timeout_secs": 45, "workspace_prefix": "channels/discord/", "emit_rate_limit": { @@ -55,8 +55,12 @@ }, "config": { "require_signature_verification": true, + "webhook_secret": null, + "polling_enabled": false, + "poll_interval_ms": 30000, + "mention_channel_ids": [], "owner_id": null, "dm_policy": "pairing", "allow_from": [] } -} \ No newline at end of file +} diff --git a/channels-src/discord/src/lib.rs b/channels-src/discord/src/lib.rs index c8b37428..cdb6c515 100644 --- a/channels-src/discord/src/lib.rs +++ b/channels-src/discord/src/lib.rs @@ -14,7 +14,7 @@ //! //! # Security //! -//! - Signature validation is handled by the host (webhook secrets) +//! - Signature validation is handled in-channel using Discord's Ed25519 headers //! - Bot token is injected by host during HTTP requests //! - WASM never sees raw credentials @@ -23,11 +23,14 @@ wit_bindgen::generate!({ path: "../../wit/channel.wit", }); +use std::{cmp::Ordering, collections::HashMap}; + +use ed25519_dalek::{Signature, Verifier, VerifyingKey}; use serde::{Deserialize, Serialize}; use exports::near::agent::channel::{ AgentResponse, ChannelConfig, Guest, HttpEndpointConfig, IncomingHttpRequest, - OutgoingHttpResponse, StatusUpdate, + OutgoingHttpResponse, PollConfig, StatusUpdate, }; use near::agent::channel_host::{self, EmittedMessage}; @@ -105,23 +108,70 @@ struct DiscordMessage { author: DiscordUser, } -/// Metadata stored with emitted messages for response routing. -#[derive(Debug, Serialize, Deserialize)] -struct DiscordMessageMetadata { - /// Discord channel ID +#[derive(Debug, Deserialize)] +struct DiscordChannelMessage { + id: String, + content: String, channel_id: String, + author: DiscordChannelAuthor, + #[serde(default)] + mentions: Vec, + #[serde(default)] + webhook_id: Option, +} - /// Interaction ID for followups - interaction_id: String, +#[derive(Debug, Deserialize)] +struct DiscordChannelAuthor { + id: String, + username: String, + global_name: Option, + #[serde(default)] + bot: bool, +} - /// Interaction token for responding - token: String, +#[derive(Debug, Clone, Serialize, Deserialize)] +struct DiscordRuntimeConfig { + #[serde(default = "default_require_signature_verification")] + require_signature_verification: bool, + #[serde(default)] + webhook_secret: Option, + #[serde(default)] + polling_enabled: bool, + #[serde(default = "default_poll_interval_ms")] + poll_interval_ms: u32, + #[serde(default)] + mention_channel_ids: Vec, + #[serde(default)] + owner_id: Option, + #[serde(default = "default_dm_policy")] + dm_policy: String, + #[serde(default)] + allow_from: Vec, +} - /// Application ID - application_id: String, +fn default_poll_interval_ms() -> u32 { + 30_000 +} - /// Thread ID (for forum threads) - thread_id: Option, +fn default_require_signature_verification() -> bool { + true +} + +fn default_dm_policy() -> String { + "pairing".to_string() +} + +fn default_runtime_config() -> DiscordRuntimeConfig { + DiscordRuntimeConfig { + require_signature_verification: default_require_signature_verification(), + webhook_secret: None, + polling_enabled: false, + poll_interval_ms: default_poll_interval_ms(), + mention_channel_ids: Vec::new(), + owner_id: None, + dm_policy: default_dm_policy(), + allow_from: Vec::new(), + } } /// Workspace path for persisting owner_id across WASM callbacks. @@ -133,30 +183,71 @@ const ALLOW_FROM_PATH: &str = "state/allow_from"; /// Channel name for pairing store (used by pairing host APIs). const CHANNEL_NAME: &str = "discord"; -/// Channel configuration from capabilities file. -#[derive(Debug, Deserialize)] -struct DiscordConfig { +/// Metadata stored with emitted messages for response routing. +#[derive(Debug, Serialize, Deserialize)] +struct DiscordMessageMetadata { + /// Discord channel ID + channel_id: String, + + /// Interaction ID for followups #[serde(default)] - #[allow(dead_code)] - require_signature_verification: bool, + interaction_id: Option, + + /// Interaction token for responding #[serde(default)] - owner_id: Option, + token: Option, + + /// Application ID #[serde(default)] - dm_policy: Option, + application_id: Option, + + /// Source message ID when handling mention-poll events. #[serde(default)] - allow_from: Option>, + source_message_id: Option, + + /// Thread ID (for forum threads) + thread_id: Option, } struct DiscordChannel; impl Guest for DiscordChannel { fn on_start(config_json: String) -> Result { - let config: DiscordConfig = serde_json::from_str(&config_json) - .map_err(|e| format!("Failed to parse config: {}", e))?; - channel_host::log(channel_host::LogLevel::Info, "Discord channel starting"); - // Persist owner_id so subsequent callbacks can read it + let config = + serde_json::from_str::(&config_json).unwrap_or_else(|e| { + channel_host::log( + channel_host::LogLevel::Warn, + &format!("Invalid config JSON, using defaults: {}", e), + ); + default_runtime_config() + }); + + if let Ok(serialized) = serde_json::to_string(&config) { + let _ = channel_host::workspace_write("config.json", &serialized); + } + + if config.require_signature_verification + && config + .webhook_secret + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + .is_none() + { + channel_host::log( + channel_host::LogLevel::Error, + "Discord channel misconfigured: require_signature_verification=true but webhook_secret is empty", + ); + } else if !config.require_signature_verification { + channel_host::log( + channel_host::LogLevel::Warn, + "Discord signature verification is disabled; webhook endpoint is unprotected", + ); + } + + // Persist owner_id so subsequent callbacks can read it. if let Some(ref owner_id) = config.owner_id { let _ = channel_host::workspace_write(OWNER_ID_PATH, owner_id); channel_host::log( @@ -167,12 +258,10 @@ impl Guest for DiscordChannel { let _ = channel_host::workspace_write(OWNER_ID_PATH, ""); } - // Persist dm_policy and allow_from for DM pairing - let dm_policy = config.dm_policy.as_deref().unwrap_or("pairing"); - let _ = channel_host::workspace_write(DM_POLICY_PATH, dm_policy); - - let allow_from_json = serde_json::to_string(&config.allow_from.unwrap_or_default()) - .unwrap_or_else(|_| "[]".to_string()); + // Persist dm_policy and allow_from for DM pairing. + let _ = channel_host::workspace_write(DM_POLICY_PATH, &config.dm_policy); + let allow_from_json = + serde_json::to_string(&config.allow_from).unwrap_or_else(|_| "[]".to_string()); let _ = channel_host::workspace_write(ALLOW_FROM_PATH, &allow_from_json); Ok(ChannelConfig { @@ -180,13 +269,59 @@ impl Guest for DiscordChannel { http_endpoints: vec![HttpEndpointConfig { path: "/webhook/discord".to_string(), methods: vec!["POST".to_string()], - require_secret: true, + require_secret: false, }], - poll: None, + poll: if config.polling_enabled { + Some(PollConfig { + interval_ms: config.poll_interval_ms.max(30_000), + enabled: true, + }) + } else { + None + }, }) } fn on_http_request(req: IncomingHttpRequest) -> OutgoingHttpResponse { + let config = load_runtime_config(); + let headers: HashMap = + serde_json::from_str(&req.headers_json).unwrap_or_default(); + if config.require_signature_verification { + if config + .webhook_secret + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + .is_none() + { + channel_host::log( + channel_host::LogLevel::Error, + "Discord channel misconfigured: webhook_secret not set while verification is required", + ); + return json_response( + 500, + serde_json::json!({"error": "Channel misconfigured: webhook_secret not set"}), + ); + } + + if !verify_discord_request_signature( + headers, + &req.body, + config.webhook_secret.as_deref(), + ) { + channel_host::log( + channel_host::LogLevel::Warn, + "Discord signature verification failed", + ); + return json_response(401, serde_json::json!({"error": "Invalid signature"})); + } + } else { + channel_host::log( + channel_host::LogLevel::Warn, + "Discord signature verification is disabled; accepting unverified webhook request", + ); + } + let body_str = match std::str::from_utf8(&req.body) { Ok(s) => s, Err(_) => { @@ -215,9 +350,16 @@ impl Guest for DiscordChannel { // Application Command (slash command) 2 => { if handle_slash_command(&interaction) { - json_response(200, serde_json::json!({"type": 5})) + json_response( + 200, + serde_json::json!({ + "type": 5, + "data": { + "content": "🤔 Thinking..." + } + }), + ) } else { - // Permission denied — ephemeral response json_response( 200, serde_json::json!({ @@ -252,24 +394,18 @@ impl Guest for DiscordChannel { } } - fn on_poll() {} + fn on_poll() { + poll_for_mentions(); + } fn on_respond(response: AgentResponse) -> Result<(), String> { let metadata: DiscordMessageMetadata = serde_json::from_str(&response.metadata_json) .map_err(|e| format!("Failed to parse metadata: {}", e))?; - // Use webhook endpoint for followup - let url = format!( - "https://discord.com/api/v10/webhooks/{}/{}", - metadata.application_id, metadata.token - ); - // Truncate content to 2000 characters to comply with Discord limits let content = truncate_message(&response.content); - let mut payload = serde_json::json!({ - "content": content, - }); + let mut payload = serde_json::json!({ "content": content }); // Check for embeds in metadata if let Ok(meta_json) = serde_json::from_str::(&response.metadata_json) { @@ -285,29 +421,50 @@ impl Guest for DiscordChannel { "Content-Type": "application/json" }); + let (method, url) = if let (Some(application_id), Some(token)) = + (metadata.application_id.as_ref(), metadata.token.as_ref()) + { + ( + "PATCH", + format!( + "https://discord.com/api/v10/webhooks/{}/{}/messages/@original", + application_id, token + ), + ) + } else if let Some(source_message_id) = metadata.source_message_id.as_ref() { + payload["message_reference"] = serde_json::json!({ + "message_id": source_message_id + }); + payload["allowed_mentions"] = serde_json::json!({ + "replied_user": true + }); + let mention_payload = serde_json::to_vec(&payload) + .map_err(|e| format!("Failed to serialize mention payload: {}", e))?; + let mention_url = format!( + "https://discord.com/api/v10/channels/{}/messages", + metadata.channel_id + ); + let result = channel_host::http_request( + "POST", + &mention_url, + &discord_auth_headers_json(true), + Some(&mention_payload), + None, + ); + return map_discord_response(result); + } else { + return Err("Unsupported Discord response metadata".to_string()); + }; + let result = channel_host::http_request( - "POST", + method, &url, &headers.to_string(), Some(&payload_bytes), None, ); - match result { - Ok(http_response) => { - if http_response.status >= 200 && http_response.status < 300 { - channel_host::log(channel_host::LogLevel::Debug, "Posted followup to Discord"); - Ok(()) - } else { - let body_str = String::from_utf8_lossy(&http_response.body); - Err(format!( - "Discord API error: {} - {}", - http_response.status, body_str - )) - } - } - Err(e) => Err(format!("HTTP request failed: {}", e)), - } + map_discord_response(result) } fn on_status(_update: StatusUpdate) {} @@ -324,7 +481,442 @@ impl Guest for DiscordChannel { } } -/// Returns true if the message was emitted, false if permission denied. +fn map_discord_response( + result: Result, +) -> Result<(), String> { + match result { + Ok(http_response) => { + if http_response.status >= 200 && http_response.status < 300 { + channel_host::log(channel_host::LogLevel::Debug, "Posted response to Discord"); + Ok(()) + } else { + let body_str = String::from_utf8_lossy(&http_response.body); + Err(format!( + "Discord API error: {} - {}", + http_response.status, body_str + )) + } + } + Err(e) => Err(format!("HTTP request failed: {}", e)), + } +} + +fn load_runtime_config() -> DiscordRuntimeConfig { + channel_host::workspace_read("config.json") + .and_then(|raw| serde_json::from_str::(&raw).ok()) + .unwrap_or_else(default_runtime_config) +} + +fn poll_for_mentions() { + let config = load_runtime_config(); + if !config.polling_enabled || config.mention_channel_ids.is_empty() { + return; + } + + let bot_id = match get_or_fetch_bot_id() { + Some(id) => id, + None => { + channel_host::log( + channel_host::LogLevel::Warn, + "Skipping mention polling: failed to resolve bot user id", + ); + return; + } + }; + + for channel_id in &config.mention_channel_ids { + poll_channel_mentions(channel_id, &bot_id); + } +} + +fn get_or_fetch_bot_id() -> Option { + if let Some(id) = channel_host::workspace_read("bot_user_id.txt") { + let trimmed = id.trim(); + if !trimmed.is_empty() { + return Some(trimmed.to_string()); + } + } + + let response = channel_host::http_request( + "GET", + "https://discord.com/api/v10/users/@me", + &discord_auth_headers_json(false), + None, + Some(10_000), + ) + .ok()?; + + if !(200..300).contains(&response.status) { + return None; + } + + let value: serde_json::Value = serde_json::from_slice(&response.body).ok()?; + let id = value.get("id")?.as_str()?.to_string(); + let _ = channel_host::workspace_write("bot_user_id.txt", &id); + Some(id) +} + +fn poll_channel_mentions(channel_id: &str, bot_id: &str) { + let cursor_path = format!("cursor_{}.txt", channel_id); + let last_seen = channel_host::workspace_read(&cursor_path).map(|s| s.trim().to_string()); + + // On first run for a channel, initialize the cursor to "latest seen" and + // skip back-processing historical messages. + if last_seen.is_none() { + if let Some(latest) = fetch_latest_message_id(channel_id) { + let _ = channel_host::workspace_write(&cursor_path, &latest); + } + return; + } + + let Some(mut messages) = + fetch_messages_after_cursor(channel_id, last_seen.as_deref().unwrap_or("")) + else { + return; + }; + if messages.is_empty() { + return; + } + + messages.sort_by(|a, b| compare_message_ids(&a.id, &b.id)); + let mut max_seen = last_seen.clone(); + let mut recent_ids = load_recent_processed_ids(channel_id); + let mut dedup_updated = false; + + for msg in messages { + if is_new_message(max_seen.as_deref(), &msg.id) { + max_seen = Some(msg.id.clone()); + } + + if msg.webhook_id.is_some() || msg.author.bot || msg.author.id == bot_id { + continue; + } + + if !message_mentions_bot(&msg, bot_id) { + continue; + } + + if recent_ids.iter().any(|id| id == &msg.id) { + continue; + } + + let user_name = msg + .author + .global_name + .as_ref() + .filter(|s| !s.is_empty()) + .unwrap_or(&msg.author.username) + .clone(); + if !check_sender_permission(&msg.author.id, Some(&user_name), false, None) { + continue; + } + + let content = strip_bot_mention(&msg.content, bot_id); + let metadata = DiscordMessageMetadata { + channel_id: msg.channel_id.clone(), + interaction_id: None, + token: None, + application_id: None, + source_message_id: Some(msg.id.clone()), + thread_id: None, + }; + + let metadata_json = match serde_json::to_string(&metadata) { + Ok(v) => v, + Err(e) => { + channel_host::log( + channel_host::LogLevel::Warn, + &format!("Failed to serialize mention metadata: {}", e), + ); + continue; + } + }; + + channel_host::emit_message(&EmittedMessage { + user_id: msg.author.id.clone(), + user_name: Some(user_name.clone()), + content: if content.is_empty() { + "mention".to_string() + } else { + content + }, + thread_id: None, + metadata_json, + attachments: vec![], + }); + + remember_processed_id(&mut recent_ids, &msg.id); + dedup_updated = true; + } + + if let Some(cursor) = max_seen { + let _ = channel_host::workspace_write(&cursor_path, &cursor); + } + if dedup_updated { + let _ = save_recent_processed_ids(channel_id, &recent_ids); + } +} + +fn fetch_latest_message_id(channel_id: &str) -> Option { + let url = format!( + "https://discord.com/api/v10/channels/{}/messages?limit=1", + channel_id + ); + let response = channel_host::http_request( + "GET", + &url, + &discord_auth_headers_json(false), + None, + Some(10_000), + ) + .ok()?; + if !(200..300).contains(&response.status) { + let body = String::from_utf8_lossy(&response.body); + channel_host::log( + channel_host::LogLevel::Warn, + &format!( + "Discord initial poll failed for channel {}: status={} body={}", + channel_id, response.status, body + ), + ); + return None; + } + let messages: Vec = serde_json::from_slice(&response.body).ok()?; + messages.first().map(|m| m.id.clone()) +} + +fn fetch_messages_after_cursor( + channel_id: &str, + last_seen: &str, +) -> Option> { + const PAGE_LIMIT: usize = 100; + const MAX_PAGES: usize = 50; + + let mut all_messages = Vec::new(); + let mut after = last_seen.to_string(); + + for page in 0..MAX_PAGES { + let url = format!( + "https://discord.com/api/v10/channels/{}/messages?limit={}&after={}", + channel_id, PAGE_LIMIT, after + ); + let response = match channel_host::http_request( + "GET", + &url, + &discord_auth_headers_json(false), + None, + Some(10_000), + ) { + Ok(r) => r, + Err(e) => { + channel_host::log( + channel_host::LogLevel::Warn, + &format!( + "Discord poll request failed for channel {}: {}", + channel_id, e + ), + ); + return None; + } + }; + + if !(200..300).contains(&response.status) { + let body = String::from_utf8_lossy(&response.body); + channel_host::log( + channel_host::LogLevel::Warn, + &format!( + "Discord poll failed for channel {}: status={} body={}", + channel_id, response.status, body + ), + ); + return None; + } + + let messages: Vec = match serde_json::from_slice(&response.body) { + Ok(v) => v, + Err(e) => { + channel_host::log( + channel_host::LogLevel::Warn, + &format!("Failed to parse polled Discord messages: {}", e), + ); + return None; + } + }; + let page_len = messages.len(); + if messages.is_empty() { + break; + } + + let page_max_id = messages + .iter() + .map(|m| m.id.as_str()) + .max_by(|a, b| compare_message_ids(a, b)) + .map(str::to_string); + + all_messages.extend(messages.into_iter()); + + if page_len < PAGE_LIMIT { + break; + } + + if let Some(max_id) = page_max_id { + if max_id == after { + break; + } + after = max_id; + } else { + break; + } + + if page + 1 == MAX_PAGES { + channel_host::log( + channel_host::LogLevel::Warn, + &format!( + "Discord poll pagination limit reached for channel {}; processing partial batch", + channel_id + ), + ); + } + } + + Some(all_messages) +} + +fn compare_message_ids(a: &str, b: &str) -> Ordering { + match (a.parse::(), b.parse::()) { + (Ok(left), Ok(right)) => left.cmp(&right), + _ => a.cmp(b), + } +} + +fn dedup_ids_path(channel_id: &str) -> String { + format!("dedup_{}.json", channel_id) +} + +fn load_recent_processed_ids(channel_id: &str) -> Vec { + let path = dedup_ids_path(channel_id); + channel_host::workspace_read(&path) + .and_then(|raw| serde_json::from_str::>(&raw).ok()) + .unwrap_or_default() +} + +fn save_recent_processed_ids(channel_id: &str, ids: &[String]) -> Result<(), String> { + let path = dedup_ids_path(channel_id); + let raw = + serde_json::to_string(ids).map_err(|e| format!("Failed to serialize dedup ids: {}", e))?; + channel_host::workspace_write(&path, &raw) +} + +fn remember_processed_id(ids: &mut Vec, message_id: &str) { + const MAX_RECENT_IDS: usize = 200; + if ids.iter().any(|id| id == message_id) { + return; + } + ids.push(message_id.to_string()); + if ids.len() > MAX_RECENT_IDS { + let drop_count = ids.len() - MAX_RECENT_IDS; + ids.drain(0..drop_count); + } +} + +fn is_new_message(last_seen: Option<&str>, current: &str) -> bool { + match last_seen { + None => true, + Some(prev) => { + let prev_num = prev.parse::().ok(); + let cur_num = current.parse::().ok(); + match (prev_num, cur_num) { + (Some(p), Some(c)) => c > p, + _ => current > prev, + } + } + } +} + +fn message_mentions_bot(msg: &DiscordChannelMessage, bot_id: &str) -> bool { + msg.mentions.iter().any(|u| u.id == bot_id) + || msg.content.contains(&format!("<@{}>", bot_id)) + || msg.content.contains(&format!("<@!{}>", bot_id)) +} + +fn strip_bot_mention(content: &str, bot_id: &str) -> String { + content + .replace(&format!("<@{}>", bot_id), "") + .replace(&format!("<@!{}>", bot_id), "") + .trim() + .to_string() +} + +fn discord_auth_headers_json(include_content_type: bool) -> String { + if include_content_type { + serde_json::json!({ + "Content-Type": "application/json", + "Authorization": "Bot {DISCORD_BOT_TOKEN}" + }) + .to_string() + } else { + serde_json::json!({ + "Authorization": "Bot {DISCORD_BOT_TOKEN}" + }) + .to_string() + } +} + +fn verify_discord_request_signature( + headers: HashMap, + body: &[u8], + public_key_hex: Option<&str>, +) -> bool { + let Some(public_key_hex) = public_key_hex.map(str::trim).filter(|s| !s.is_empty()) else { + return false; + }; + let Some(signature_hex) = header_case_insensitive(&headers, "x-signature-ed25519") else { + return false; + }; + let Some(timestamp) = header_case_insensitive(&headers, "x-signature-timestamp") else { + return false; + }; + + let public_key_bytes = match hex::decode(public_key_hex) { + Ok(v) => v, + Err(_) => return false, + }; + let public_key_arr: [u8; 32] = match public_key_bytes.try_into() { + Ok(v) => v, + Err(_) => return false, + }; + let verifying_key = match VerifyingKey::from_bytes(&public_key_arr) { + Ok(v) => v, + Err(_) => return false, + }; + + let sig_bytes = match hex::decode(signature_hex.trim()) { + Ok(v) => v, + Err(_) => return false, + }; + let sig_arr: [u8; 64] = match sig_bytes.try_into() { + Ok(v) => v, + Err(_) => return false, + }; + let signature = Signature::from_bytes(&sig_arr); + + let mut signed_message = Vec::with_capacity(timestamp.len() + body.len()); + signed_message.extend_from_slice(timestamp.as_bytes()); + signed_message.extend_from_slice(body); + + verifying_key.verify(&signed_message, &signature).is_ok() +} + +fn header_case_insensitive<'a>( + headers: &'a HashMap, + name: &str, +) -> Option<&'a str> { + headers + .iter() + .find(|(k, _)| k.eq_ignore_ascii_case(name)) + .map(|(_, v)| v.as_str()) +} + fn handle_slash_command(interaction: &DiscordInteraction) -> bool { let user = interaction .member @@ -342,10 +934,8 @@ fn handle_slash_command(interaction: &DiscordInteraction) -> bool { }) .unwrap_or_default(); - // DM if no guild member context (only direct user field set) + // DM if no guild member context (only direct user field set). let is_dm = interaction.member.is_none(); - - // Permission check if !check_sender_permission( &user_id, Some(&user_name), @@ -380,9 +970,10 @@ fn handle_slash_command(interaction: &DiscordInteraction) -> bool { let metadata = DiscordMessageMetadata { channel_id: channel_id.clone(), - interaction_id: interaction.id.clone(), - token: interaction.token.clone(), - application_id: interaction.application_id.clone(), + interaction_id: Some(interaction.id.clone()), + token: Some(interaction.token.clone()), + application_id: Some(interaction.application_id.clone()), + source_message_id: None, thread_id: None, }; @@ -393,13 +984,14 @@ fn handle_slash_command(interaction: &DiscordInteraction) -> bool { channel_host::LogLevel::Error, &format!("Failed to serialize metadata: {}", e), ); + // Attempt to notify user of internal error let url = format!( "https://discord.com/api/v10/webhooks/{}/{}", interaction.application_id, interaction.token ); let payload = serde_json::json!({ "content": "❌ Internal Error: Failed to process command metadata.", - "flags": 64 + "flags": 64 // Ephemeral }); let _ = channel_host::http_request( "POST", @@ -408,7 +1000,7 @@ fn handle_slash_command(interaction: &DiscordInteraction) -> bool { Some(&serde_json::to_vec(&payload).unwrap_or_default()), None, ); - return true; // Error, but not a permission denial + return true; } }; @@ -424,6 +1016,7 @@ fn handle_slash_command(interaction: &DiscordInteraction) -> bool { } fn handle_message_component(interaction: &DiscordInteraction, message: &DiscordMessage) { + // Check member first (for server contexts), then user (for DMs) let user = interaction .member .as_ref() @@ -449,9 +1042,10 @@ fn handle_message_component(interaction: &DiscordInteraction, message: &DiscordM let metadata = DiscordMessageMetadata { channel_id: channel_id.clone(), - interaction_id: interaction.id.clone(), - token: interaction.token.clone(), - application_id: interaction.application_id.clone(), + interaction_id: Some(interaction.id.clone()), + token: Some(interaction.token.clone()), + application_id: Some(interaction.application_id.clone()), + source_message_id: None, thread_id: None, }; @@ -476,10 +1070,6 @@ fn handle_message_component(interaction: &DiscordInteraction, message: &DiscordM }); } -// ============================================================================ -// Permission & Pairing -// ============================================================================ - /// Context needed to send a pairing reply via Discord webhook followup. struct PairingReplyCtx { application_id: String, @@ -494,7 +1084,7 @@ fn check_sender_permission( is_dm: bool, reply_ctx: Option<&PairingReplyCtx>, ) -> bool { - // 1. Owner check (highest priority, applies to all contexts) + // 1. Owner check (highest priority, applies to all contexts). let owner_id = channel_host::workspace_read(OWNER_ID_PATH).filter(|s| !s.is_empty()); if let Some(ref owner) = owner_id { if user_id != owner { @@ -510,28 +1100,26 @@ fn check_sender_permission( return true; } - // 2. DM policy (only for DMs when no owner_id) + // 2. DM policy (only for DMs when no owner_id). if !is_dm { - return true; // Guild interactions bypass DM policy + return true; } let dm_policy = - channel_host::workspace_read(DM_POLICY_PATH).unwrap_or_else(|| "pairing".to_string()); - + channel_host::workspace_read(DM_POLICY_PATH).unwrap_or_else(|| default_dm_policy()); if dm_policy == "open" { return true; } - // 3. Build merged allow list: config allow_from + pairing store + // 3. Build merged allow list: config allow_from + pairing store. let mut allowed: Vec = channel_host::workspace_read(ALLOW_FROM_PATH) .and_then(|s| serde_json::from_str(&s).ok()) .unwrap_or_default(); - if let Ok(store_allowed) = channel_host::pairing_read_allow_from(CHANNEL_NAME) { allowed.extend(store_allowed); } - // 4. Check sender against allow list + // 4. Check sender against allow list. let is_allowed = allowed.contains(&"*".to_string()) || allowed.contains(&user_id.to_string()) || username.is_some_and(|u| allowed.contains(&u.to_string())); @@ -540,22 +1128,18 @@ fn check_sender_permission( return true; } - // 5. Not allowed — handle by policy + // 5. Not allowed - handle by policy. if dm_policy == "pairing" { let meta = serde_json::json!({ "user_id": user_id, "username": username, }) .to_string(); - match channel_host::pairing_upsert_request(CHANNEL_NAME, user_id, &meta) { Ok(result) => { channel_host::log( channel_host::LogLevel::Info, - &format!( - "Pairing request for user {}: code {}", - user_id, result.code - ), + &format!("Pairing request for user {}: code {}", user_id, result.code), ); if result.created { if let Some(ctx) = reply_ctx { @@ -580,20 +1164,16 @@ fn send_pairing_reply(ctx: &PairingReplyCtx, code: &str) -> Result<(), String> { "https://discord.com/api/v10/webhooks/{}/{}", ctx.application_id, ctx.token ); - let payload = serde_json::json!({ "content": format!( "To pair with this bot, run: `ironclaw pairing approve discord {}`", code ), - "flags": 64 // Ephemeral — only visible to the sender + "flags": 64 }); - let payload_bytes = serde_json::to_vec(&payload).map_err(|e| format!("Failed to serialize: {}", e))?; - let headers = serde_json::json!({"Content-Type": "application/json"}); - let result = channel_host::http_request( "POST", &url, @@ -601,7 +1181,6 @@ fn send_pairing_reply(ctx: &PairingReplyCtx, code: &str) -> Result<(), String> { Some(&payload_bytes), None, ); - match result { Ok(response) if response.status >= 200 && response.status < 300 => Ok(()), Ok(response) => { @@ -648,6 +1227,7 @@ fn truncate_message(content: &str) -> String { #[cfg(test)] mod tests { use super::*; + use ed25519_dalek::{Signer, SigningKey}; #[test] fn test_truncate_message() { @@ -679,15 +1259,309 @@ mod tests { fn test_metadata_serialization() { let metadata = DiscordMessageMetadata { channel_id: "123".into(), - interaction_id: "456".into(), - token: "abc".into(), - application_id: "789".into(), + interaction_id: Some("456".into()), + token: Some("abc".into()), + application_id: Some("789".into()), + source_message_id: None, thread_id: None, }; let json = serde_json::to_string(&metadata).unwrap(); let parsed: DiscordMessageMetadata = serde_json::from_str(&json).unwrap(); assert_eq!(parsed.channel_id, "123"); - assert_eq!(parsed.interaction_id, "456"); + assert_eq!(parsed.interaction_id.as_deref(), Some("456")); + } + + #[test] + fn test_is_new_message() { + assert!(is_new_message(None, "100")); + assert!(is_new_message(Some("100"), "200")); + assert!(!is_new_message(Some("200"), "100")); + assert!(!is_new_message(Some("100"), "100")); + assert!(is_new_message(Some("abc"), "abd")); + assert!(!is_new_message(Some("abd"), "abc")); + } + + #[test] + fn test_strip_bot_mention() { + assert_eq!(strip_bot_mention("<@123> hello", "123"), "hello"); + assert_eq!(strip_bot_mention("<@!123> hello", "123"), "hello"); + assert_eq!(strip_bot_mention("<@123>", "123"), ""); + assert_eq!( + strip_bot_mention("hello <@123> world <@!123>", "123"), + "hello world" + ); + } + + #[test] + fn test_message_mentions_bot() { + let msg = DiscordChannelMessage { + id: "1".to_string(), + content: "hello <@123>".to_string(), + channel_id: "10".to_string(), + author: DiscordChannelAuthor { + id: "u1".to_string(), + username: "alice".to_string(), + global_name: None, + bot: false, + }, + mentions: vec![], + webhook_id: None, + }; + assert!(message_mentions_bot(&msg, "123")); + assert!(!message_mentions_bot(&msg, "999")); + } + + #[test] + fn test_message_mentions_bot_via_mentions_array() { + let msg = DiscordChannelMessage { + id: "2".to_string(), + content: "hello".to_string(), + channel_id: "10".to_string(), + author: DiscordChannelAuthor { + id: "u1".to_string(), + username: "alice".to_string(), + global_name: None, + bot: false, + }, + mentions: vec![DiscordUser { + id: "777".to_string(), + username: "bot".to_string(), + global_name: None, + }], + webhook_id: None, + }; + assert!(message_mentions_bot(&msg, "777")); + } + + #[test] + fn test_compare_message_ids_numeric_and_lexical_fallback() { + assert_eq!(compare_message_ids("100", "20"), Ordering::Greater); + assert_eq!(compare_message_ids("20", "100"), Ordering::Less); + assert_eq!(compare_message_ids("abc", "abd"), Ordering::Less); + assert_eq!(compare_message_ids("abd", "abc"), Ordering::Greater); + } + + #[test] + fn test_remember_processed_id_dedup_and_cap() { + let mut ids = Vec::new(); + for i in 0..220 { + remember_processed_id(&mut ids, &format!("{}", i)); + } + assert_eq!(ids.len(), 200); + assert_eq!(ids.first().map(String::as_str), Some("20")); + assert_eq!(ids.last().map(String::as_str), Some("219")); + + remember_processed_id(&mut ids, "219"); + assert_eq!(ids.len(), 200); + assert_eq!(ids.last().map(String::as_str), Some("219")); + } + + #[test] + fn test_header_case_insensitive() { + let mut headers = HashMap::new(); + headers.insert("X-Signature-Timestamp".to_string(), "123".to_string()); + assert_eq!( + header_case_insensitive(&headers, "x-signature-timestamp"), + Some("123") + ); + assert_eq!(header_case_insensitive(&headers, "missing"), None); + } + + #[test] + fn test_discord_auth_headers_json_shape() { + let with_ct: serde_json::Value = + serde_json::from_str(&discord_auth_headers_json(true)).unwrap(); + assert_eq!( + with_ct.get("Content-Type").and_then(|v| v.as_str()), + Some("application/json") + ); + assert_eq!( + with_ct.get("Authorization").and_then(|v| v.as_str()), + Some("Bot {DISCORD_BOT_TOKEN}") + ); + + let no_ct: serde_json::Value = + serde_json::from_str(&discord_auth_headers_json(false)).unwrap(); + assert!(no_ct.get("Content-Type").is_none()); + assert_eq!( + no_ct.get("Authorization").and_then(|v| v.as_str()), + Some("Bot {DISCORD_BOT_TOKEN}") + ); + } + + #[test] + fn test_verify_discord_request_signature_valid() { + let signing_key = SigningKey::from_bytes(&[7u8; 32]); + let public_key_hex = hex::encode(signing_key.verifying_key().to_bytes()); + let timestamp = "1234567890"; + let body = br#"{"type":1}"#; + + let mut signed = Vec::new(); + signed.extend_from_slice(timestamp.as_bytes()); + signed.extend_from_slice(body); + let signature = signing_key.sign(&signed); + + let mut headers = HashMap::new(); + headers.insert( + "x-signature-ed25519".to_string(), + hex::encode(signature.to_bytes()), + ); + headers.insert("x-signature-timestamp".to_string(), timestamp.to_string()); + + assert!(verify_discord_request_signature( + headers, + body, + Some(&public_key_hex) + )); + } + + #[test] + fn test_verify_discord_request_signature_tampered_body() { + let signing_key = SigningKey::from_bytes(&[9u8; 32]); + let public_key_hex = hex::encode(signing_key.verifying_key().to_bytes()); + let timestamp = "1234567890"; + let body = b"hello"; + + let mut signed = Vec::new(); + signed.extend_from_slice(timestamp.as_bytes()); + signed.extend_from_slice(body); + let signature = signing_key.sign(&signed); + + let mut headers = HashMap::new(); + headers.insert( + "x-signature-ed25519".to_string(), + hex::encode(signature.to_bytes()), + ); + headers.insert("x-signature-timestamp".to_string(), timestamp.to_string()); + + assert!(!verify_discord_request_signature( + headers, + b"hello-modified", + Some(&public_key_hex) + )); + } + + #[test] + fn test_verify_discord_request_signature_wrong_public_key() { + let signing_key = SigningKey::from_bytes(&[11u8; 32]); + let wrong_key = SigningKey::from_bytes(&[12u8; 32]); + let timestamp = "1234567890"; + let body = b"payload"; + + let mut signed = Vec::new(); + signed.extend_from_slice(timestamp.as_bytes()); + signed.extend_from_slice(body); + let signature = signing_key.sign(&signed); + + let mut headers = HashMap::new(); + headers.insert( + "x-signature-ed25519".to_string(), + hex::encode(signature.to_bytes()), + ); + headers.insert("x-signature-timestamp".to_string(), timestamp.to_string()); + + assert!(!verify_discord_request_signature( + headers, + body, + Some(&hex::encode(wrong_key.verifying_key().to_bytes())) + )); + } + + #[test] + fn test_verify_discord_request_signature_missing_headers() { + let headers = HashMap::new(); + assert!(!verify_discord_request_signature( + headers, + b"abc", + Some("00112233445566778899aabbccddeeff00112233445566778899aabbccddeeff") + )); + } + + #[test] + fn test_verify_discord_request_signature_invalid_signature_hex() { + let mut headers = HashMap::new(); + headers.insert("x-signature-ed25519".to_string(), "not-hex".to_string()); + headers.insert( + "x-signature-timestamp".to_string(), + "1234567890".to_string(), + ); + assert!(!verify_discord_request_signature( + headers, + b"abc", + Some("00112233445566778899aabbccddeeff00112233445566778899aabbccddeeff") + )); + } + + #[test] + fn test_verify_discord_request_signature_invalid_public_key_hex() { + let mut headers = HashMap::new(); + headers.insert("x-signature-ed25519".to_string(), "00".repeat(64)); + headers.insert( + "x-signature-timestamp".to_string(), + "1234567890".to_string(), + ); + assert!(!verify_discord_request_signature( + headers, + b"abc", + Some("not-hex") + )); + } + + #[test] + fn test_verify_discord_request_signature_invalid_lengths() { + let mut headers = HashMap::new(); + headers.insert("x-signature-ed25519".to_string(), "00".repeat(10)); + headers.insert( + "x-signature-timestamp".to_string(), + "1234567890".to_string(), + ); + assert!(!verify_discord_request_signature( + headers.clone(), + b"abc", + Some("00".repeat(31).as_str()) + )); + assert!(!verify_discord_request_signature( + headers, + b"abc", + Some("00".repeat(32).as_str()) + )); + } + + #[test] + fn test_verify_discord_request_signature_case_insensitive_headers() { + let signing_key = SigningKey::from_bytes(&[13u8; 32]); + let public_key_hex = hex::encode(signing_key.verifying_key().to_bytes()); + let timestamp = "1234567890"; + let body = b"case-header"; + + let mut signed = Vec::new(); + signed.extend_from_slice(timestamp.as_bytes()); + signed.extend_from_slice(body); + let signature = signing_key.sign(&signed); + + let mut headers = HashMap::new(); + headers.insert( + "X-Signature-Ed25519".to_string(), + hex::encode(signature.to_bytes()), + ); + headers.insert("X-Signature-Timestamp".to_string(), timestamp.to_string()); + + assert!(verify_discord_request_signature( + headers, + body, + Some(&public_key_hex) + )); + } + + #[test] + fn test_verify_discord_request_signature_empty_public_key() { + let mut headers = HashMap::new(); + headers.insert("x-signature-ed25519".to_string(), "00".repeat(64)); + headers.insert( + "x-signature-timestamp".to_string(), + "1234567890".to_string(), + ); + assert!(!verify_discord_request_signature(headers, b"abc", Some(""))); } #[test] diff --git a/docs/LLM_PROVIDERS.md b/docs/LLM_PROVIDERS.md index 60ac2bbc..a581a56b 100644 --- a/docs/LLM_PROVIDERS.md +++ b/docs/LLM_PROVIDERS.md @@ -15,6 +15,7 @@ configurations. | io.net | `ionet` | `IONET_API_KEY` | Intelligence API | | Mistral | `mistral` | `MISTRAL_API_KEY` | Mistral models | | Yandex AI Studio | `yandex` | `YANDEX_API_KEY` | YandexGPT models | +| MiniMax | `minimax` | `MINIMAX_API_KEY` | MiniMax-M2.5 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 | @@ -74,6 +75,25 @@ Pull a model first: `ollama pull llama3.2` --- +## MiniMax + +[MiniMax](https://platform.minimax.io) provides high-performance language models with 204,800 token context windows. + +```env +LLM_BACKEND=minimax +MINIMAX_API_KEY=... +``` + +Available models: `MiniMax-M2.5` (default), `MiniMax-M2.5-highspeed` + +To use the China mainland endpoint, set: + +```env +MINIMAX_BASE_URL=https://api.minimaxi.com/v1 +``` + +--- + ## AWS Bedrock (requires `--features bedrock`) Uses the native AWS Converse API via `aws-sdk-bedrockruntime`. Supports standard AWS diff --git a/providers.json b/providers.json index 3b98af30..12723a6f 100644 --- a/providers.json +++ b/providers.json @@ -382,6 +382,27 @@ "can_list_models": true } }, + { + "id": "minimax", + "aliases": [ + "mini_max" + ], + "protocol": "open_ai_completions", + "default_base_url": "https://api.minimax.io/v1", + "api_key_env": "MINIMAX_API_KEY", + "api_key_required": true, + "base_url_env": "MINIMAX_BASE_URL", + "model_env": "MINIMAX_MODEL", + "default_model": "MiniMax-M2.5", + "description": "MiniMax API (MiniMax-M2.5 and MiniMax-M2.5-highspeed models)", + "setup": { + "kind": "api_key", + "secret_name": "llm_minimax_api_key", + "key_url": "https://platform.minimax.io", + "display_name": "MiniMax", + "can_list_models": false + } + }, { "id": "cloudflare", "aliases": [ diff --git a/registry/channels/discord.json b/registry/channels/discord.json index cf057245..6f5cd4e7 100644 --- a/registry/channels/discord.json +++ b/registry/channels/discord.json @@ -18,7 +18,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/discord-0.2.0-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/discord-0.2.0-wasm32-wasip2.tar.gz", "sha256": "efa1b9019fa33e243f8db1e1fcc732731d45836336bdd26ca19b6fe227ca8b69" } }, diff --git a/registry/channels/slack.json b/registry/channels/slack.json index 64b28e3b..e6d36604 100644 --- a/registry/channels/slack.json +++ b/registry/channels/slack.json @@ -18,7 +18,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/slack-0.2.1-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/slack-0.2.1-wasm32-wasip2.tar.gz", "sha256": "d4667e35126986509d862bc3a0088777305d8f41c75de83c1e223b42312ede48" } }, diff --git a/registry/channels/telegram.json b/registry/channels/telegram.json index 74336e41..36be1fc7 100644 --- a/registry/channels/telegram.json +++ b/registry/channels/telegram.json @@ -2,7 +2,7 @@ "name": "telegram", "display_name": "Telegram Channel", "kind": "channel", - "version": "0.2.2", + "version": "0.2.3", "wit_version": "0.3.0", "description": "Talk to your agent through a Telegram bot", "keywords": [ @@ -18,7 +18,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/telegram-0.2.2-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/telegram-0.2.3-wasm32-wasip2.tar.gz", "sha256": "b9a83d5a2d1285ce0ec116b354336a1f245f893291ccb01dffbcaccf89d72aed" } }, diff --git a/registry/channels/whatsapp.json b/registry/channels/whatsapp.json index d1017276..be3faf0d 100644 --- a/registry/channels/whatsapp.json +++ b/registry/channels/whatsapp.json @@ -18,7 +18,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/whatsapp-0.2.0-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/whatsapp-0.2.0-wasm32-wasip2.tar.gz", "sha256": "feb9194719d9bed796b070ab4dc30348dbfb5d3dec56f9f21e02d14137abab01" } }, diff --git a/registry/tools/github.json b/registry/tools/github.json index e2dd1168..e84f756d 100644 --- a/registry/tools/github.json +++ b/registry/tools/github.json @@ -2,7 +2,7 @@ "name": "github", "display_name": "GitHub", "kind": "tool", - "version": "0.2.1", + "version": "0.2.0", "wit_version": "0.3.0", "description": "GitHub integration for issues, PRs, repos, and code search", "keywords": [ @@ -19,7 +19,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/github-0.2.0-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/github-0.2.0-wasm32-wasip2.tar.gz", "sha256": "da9fac56b6f20197a415489bbaec9fefb085a5cf6324cab79ea48a47eb19c13b" } }, diff --git a/registry/tools/gmail.json b/registry/tools/gmail.json index dc9e6c40..08913ce6 100644 --- a/registry/tools/gmail.json +++ b/registry/tools/gmail.json @@ -18,7 +18,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/gmail-0.2.0-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/gmail-0.2.0-wasm32-wasip2.tar.gz", "sha256": "ee9574e02e92bc1d481f1310eb88afd99ee52bf6971074ab33bd76bf99b34b1d" } }, diff --git a/registry/tools/google-calendar.json b/registry/tools/google-calendar.json index 0b773f69..c43112d3 100644 --- a/registry/tools/google-calendar.json +++ b/registry/tools/google-calendar.json @@ -18,7 +18,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/google-calendar-0.2.0-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-calendar-0.2.0-wasm32-wasip2.tar.gz", "sha256": "2fa47150ea222e787c122182ad6f4dfa2ffaf5fe490d05e8de887a76445f8d2d" } }, diff --git a/registry/tools/google-docs.json b/registry/tools/google-docs.json index 66ddd407..9f1ab133 100644 --- a/registry/tools/google-docs.json +++ b/registry/tools/google-docs.json @@ -18,7 +18,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/google-docs-0.2.0-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-docs-0.2.0-wasm32-wasip2.tar.gz", "sha256": "40e134a1c1564f832ca861c3396895d4e33ec67b99313fc1f97baf8d971423a9" } }, diff --git a/registry/tools/google-drive.json b/registry/tools/google-drive.json index 6ee52089..9766e555 100644 --- a/registry/tools/google-drive.json +++ b/registry/tools/google-drive.json @@ -18,7 +18,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/google-drive-0.2.0-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-drive-0.2.0-wasm32-wasip2.tar.gz", "sha256": "002a341a1d58125563a7c69561b26fbc2629b04ea723cade744102bdc0fbb71f" } }, diff --git a/registry/tools/google-sheets.json b/registry/tools/google-sheets.json index 1cf5c808..b63265e1 100644 --- a/registry/tools/google-sheets.json +++ b/registry/tools/google-sheets.json @@ -18,7 +18,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/google-sheets-0.2.0-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-sheets-0.2.0-wasm32-wasip2.tar.gz", "sha256": "8aa2c9d52f033edea3a6c2311b0ec694ccb6d0a54ef07e94d72bf8be1ce8009a" } }, diff --git a/registry/tools/google-slides.json b/registry/tools/google-slides.json index 9c5684b8..54187531 100644 --- a/registry/tools/google-slides.json +++ b/registry/tools/google-slides.json @@ -17,7 +17,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/google-slides-0.2.0-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/google-slides-0.2.0-wasm32-wasip2.tar.gz", "sha256": "e931a97d4fd0b0b938e464dc7c7f2be6ea6b4d1508f5ea3cd931d44db23f05f5" } }, diff --git a/registry/tools/llm-context.json b/registry/tools/llm-context.json new file mode 100644 index 00000000..a647a153 --- /dev/null +++ b/registry/tools/llm-context.json @@ -0,0 +1,41 @@ +{ + "name": "llm-context", + "display_name": "LLM Context", + "kind": "tool", + "version": "0.1.0", + "wit_version": "0.3.0", + "description": "Fetch pre-extracted web content from Brave Search for grounding LLM answers (RAG, fact-checking)", + "keywords": [ + "search", + "web", + "brave", + "rag", + "grounding", + "llm", + "context" + ], + "source": { + "dir": "tools-src/llm-context", + "capabilities": "llm-context-tool.capabilities.json", + "crate_name": "llm-context-tool" + }, + "artifacts": { + "wasm32-wasip2": { + "url": "https://github.com/nearai/ironclaw/releases/latest/download/llm-context-wasm32-wasip2.tar.gz", + "sha256": "581cc5867ef3b75116b7ddc8161e63dd92befe2b53e6ad8213c007639aa243c3" + } + }, + "auth_summary": { + "method": "manual", + "provider": "Brave", + "secrets": [ + "brave_api_key" + ], + "shared_auth": "Same API key as Web Search tool (brave_api_key)", + "setup_url": "https://brave.com/search/api/" + }, + "tags": [ + "default", + "search" + ] +} diff --git a/registry/tools/slack.json b/registry/tools/slack.json index 194f1ffe..11bd7fff 100644 --- a/registry/tools/slack.json +++ b/registry/tools/slack.json @@ -17,8 +17,8 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/slack-0.2.0-wasm32-wasip2.tar.gz", - "sha256": "8af3f884240de8413d272845fad2164a347d7d2a502a0d148aa38425b93f62ed" + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/slack-0.2.1-wasm32-wasip2.tar.gz", + "sha256": "d4667e35126986509d862bc3a0088777305d8f41c75de83c1e223b42312ede48" } }, "auth_summary": { diff --git a/registry/tools/telegram.json b/registry/tools/telegram.json index 0213126b..680d6fdb 100644 --- a/registry/tools/telegram.json +++ b/registry/tools/telegram.json @@ -18,8 +18,8 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/telegram-0.2.0-wasm32-wasip2.tar.gz", - "sha256": "2c66245913854be4294021fc6bb479e43f7d65830c5cec25cf6c60a71d1af468" + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/telegram-0.2.2-wasm32-wasip2.tar.gz", + "sha256": "b9a83d5a2d1285ce0ec116b354336a1f245f893291ccb01dffbcaccf89d72aed" } }, "auth_summary": { diff --git a/registry/tools/web-search.json b/registry/tools/web-search.json index 36cc6f6b..4da5744b 100644 --- a/registry/tools/web-search.json +++ b/registry/tools/web-search.json @@ -18,7 +18,7 @@ }, "artifacts": { "wasm32-wasip2": { - "url": "https://github.com/nearai/ironclaw/releases/latest/download/web-search-0.2.0-wasm32-wasip2.tar.gz", + "url": "https://github.com/nearai/ironclaw/releases/download/v0.18.0/web-search-0.2.0-wasm32-wasip2.tar.gz", "sha256": "56834573c54ea2a33cea1eb0f04bbdf59f1ef8d8702995cf431b0921302eeccc" } }, diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index b791f6d7..d0ae98ad 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -554,6 +554,31 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { }; if needs_approval { + // In non-DM relay channels, auto-deny approval- + // requiring tools to prevent stuck AwaitingApproval + // state and prompt injection from other users. + let is_relay = self.message.channel.ends_with("-relay"); + let is_dm = self + .message + .metadata + .get("event_type") + .and_then(|v| v.as_str()) + == Some("direct_message"); + if is_relay && !is_dm { + tracing::info!( + tool = %tc.name, + channel = %self.message.channel, + "Auto-denying approval-requiring tool in non-DM relay channel" + ); + let reject_msg = format!( + "Tool '{}' requires approval and cannot run in shared channels. \ + Ask the user to message me directly (DM) to use this tool.", + tc.name + ); + preflight.push((tc, PreflightOutcome::Rejected(reject_msg))); + continue; + } + approval_needed = Some((idx, tc, tool)); break; } @@ -2235,4 +2260,51 @@ mod tests { "Present 'data' field should produce non-empty string" ); } + + /// Test the relay channel auto-deny decision logic: + /// approval-requiring tools in non-DM relay channels must be rejected. + #[test] + fn test_relay_non_dm_auto_deny_decision() { + use crate::channels::IncomingMessage; + + // Case 1: relay channel + non-DM → should auto-deny + let msg = IncomingMessage::new("slack-relay", "u1", "hello") + .with_metadata(serde_json::json!({ "event_type": "message" })); + let is_relay = msg.channel.ends_with("-relay"); + let is_dm = + msg.metadata.get("event_type").and_then(|v| v.as_str()) == Some("direct_message"); + assert!(is_relay && !is_dm, "Should auto-deny in relay non-DM"); + + // Case 2: relay channel + DM → should NOT auto-deny + let msg_dm = IncomingMessage::new("slack-relay", "u1", "hello") + .with_metadata(serde_json::json!({ "event_type": "direct_message" })); + let is_dm_2 = + msg_dm.metadata.get("event_type").and_then(|v| v.as_str()) == Some("direct_message"); + assert!( + !msg_dm.channel.ends_with("-relay") || is_dm_2, + "Should NOT auto-deny in relay DM" + ); + + // Case 3: non-relay channel → should NOT auto-deny + let msg_web = IncomingMessage::new("web", "u1", "hello") + .with_metadata(serde_json::json!({ "event_type": "message" })); + assert!( + !msg_web.channel.ends_with("-relay"), + "Non-relay channel should not trigger auto-deny" + ); + } + + /// Test that the auto-deny produces a PreflightOutcome::Rejected-style message. + #[test] + fn test_relay_auto_deny_message_format() { + let tool_name = "shell"; + let result_msg = format!( + "Tool '{}' requires approval and cannot run in shared channels. \ + Ask the user to message me directly (DM) to use this tool.", + tool_name + ); + assert!(result_msg.contains("shell")); + assert!(result_msg.contains("approval")); + assert!(result_msg.contains("DM")); + } } diff --git a/src/agent/mod.rs b/src/agent/mod.rs index de2434be..ee980233 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -32,7 +32,6 @@ pub mod task; mod thread_ops; pub mod undo; -pub use crate::worker::{Worker, WorkerDeps}; pub(crate) use agent_loop::truncate_for_preview; pub use agent_loop::{Agent, AgentDeps}; pub use compaction::{CompactionResult, ContextCompactor}; diff --git a/src/agent/routine.rs b/src/agent/routine.rs index 72226502..2dee6333 100644 --- a/src/agent/routine.rs +++ b/src/agent/routine.rs @@ -207,7 +207,7 @@ impl Trigger { #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] pub enum RoutineAction { - /// Single LLM call, no tools. Cheap and fast. + /// Single LLM call (optionally with tools). Cheap and fast. Lightweight { /// The prompt sent to the LLM. prompt: String, @@ -217,6 +217,14 @@ pub enum RoutineAction { /// Max output tokens (default: 4096). #[serde(default = "default_max_tokens")] max_tokens: u32, + /// Enable tool access (default: false for backward compatibility). + /// When true, the LLM can call tools during execution. + /// Tools requiring approval are automatically filtered out. + #[serde(default)] + use_tools: bool, + /// Max tool call rounds (default: 3). Only used when use_tools is true. + #[serde(default = "default_max_tool_rounds")] + max_tool_rounds: u32, }, /// Full multi-turn worker job with tool access. FullJob { @@ -243,6 +251,19 @@ fn default_max_iterations() -> u32 { 10 } +fn default_max_tool_rounds() -> u32 { + 3 +} + +/// Hard upper bound for max_tool_rounds to prevent runaway loops and cost explosion. +pub(crate) const MAX_TOOL_ROUNDS_LIMIT: u32 = 20; + +/// Clamp max_tool_rounds to [1, MAX_TOOL_ROUNDS_LIMIT]. +/// Accepts u64 to avoid truncation before clamping. +fn clamp_max_tool_rounds(value: u64) -> u32 { + value.clamp(1, MAX_TOOL_ROUNDS_LIMIT as u64) as u32 +} + /// Parse a `tool_permissions` JSON array into a `Vec`. pub fn parse_tool_permissions(value: &serde_json::Value) -> Vec { value @@ -290,10 +311,22 @@ impl RoutineAction { .get("max_tokens") .and_then(|v| v.as_u64()) .unwrap_or(default_max_tokens() as u64) as u32; + let use_tools = config + .get("use_tools") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let max_tool_rounds = clamp_max_tool_rounds( + config + .get("max_tool_rounds") + .and_then(|v| v.as_u64()) + .unwrap_or(default_max_tool_rounds() as u64), + ); Ok(RoutineAction::Lightweight { prompt, context_paths, max_tokens, + use_tools, + max_tool_rounds, }) } "full_job" => { @@ -339,10 +372,14 @@ impl RoutineAction { prompt, context_paths, max_tokens, + use_tools, + max_tool_rounds, } => serde_json::json!({ "prompt": prompt, "context_paths": context_paths, "max_tokens": max_tokens, + "use_tools": use_tools, + "max_tool_rounds": max_tool_rounds, }), RoutineAction::FullJob { title, @@ -504,7 +541,8 @@ pub fn next_cron_fire( #[cfg(test)] mod tests { use crate::agent::routine::{ - RoutineAction, RoutineGuardrails, RunStatus, Trigger, content_hash, next_cron_fire, + MAX_TOOL_ROUNDS_LIMIT, RoutineAction, RoutineGuardrails, RunStatus, Trigger, content_hash, + next_cron_fire, }; #[test] @@ -554,11 +592,13 @@ mod tests { prompt: "Check PRs".to_string(), context_paths: vec!["context/priorities.md".to_string()], max_tokens: 2048, + use_tools: false, + max_tool_rounds: 3, }; let json = action.to_config_json(); let parsed = RoutineAction::from_db("lightweight", json).expect("parse lightweight"); assert!( - matches!(parsed, RoutineAction::Lightweight { prompt, context_paths, max_tokens } + matches!(parsed, RoutineAction::Lightweight { prompt, context_paths, max_tokens, .. } if prompt == "Check PRs" && context_paths.len() == 1 && max_tokens == 2048) ); } @@ -695,4 +735,77 @@ mod tests { ); assert_eq!(Trigger::Manual.type_tag(), "manual"); } + + #[test] + fn test_action_lightweight_backward_compat_no_use_tools() { + // Simulate old DB record without use_tools field + let json = serde_json::json!({ + "prompt": "old routine", + "context_paths": [], + "max_tokens": 4096 + }); + let parsed = RoutineAction::from_db("lightweight", json).expect("parse lightweight"); + assert!( + matches!(parsed, RoutineAction::Lightweight { use_tools, max_tool_rounds, .. } + if !use_tools && max_tool_rounds == 3), + "missing use_tools should default to false, max_tool_rounds to 3" + ); + } + + #[test] + fn test_max_tool_rounds_clamped_to_upper_bound() { + let json = serde_json::json!({ + "prompt": "test", + "use_tools": true, + "max_tool_rounds": 9999 + }); + let parsed = RoutineAction::from_db("lightweight", json).expect("parse"); + match parsed { + RoutineAction::Lightweight { + max_tool_rounds, .. + } => { + assert_eq!( + max_tool_rounds, MAX_TOOL_ROUNDS_LIMIT, + "should clamp to MAX_TOOL_ROUNDS_LIMIT" + ); + } + _ => panic!("expected Lightweight"), + } + } + + #[test] + fn test_max_tool_rounds_clamped_to_lower_bound() { + let json = serde_json::json!({ + "prompt": "test", + "use_tools": true, + "max_tool_rounds": 0 + }); + let parsed = RoutineAction::from_db("lightweight", json).expect("parse"); + match parsed { + RoutineAction::Lightweight { + max_tool_rounds, .. + } => { + assert_eq!(max_tool_rounds, 1, "should clamp 0 to 1"); + } + _ => panic!("expected Lightweight"), + } + } + + #[test] + fn test_max_tool_rounds_normal_value_passes_through() { + let json = serde_json::json!({ + "prompt": "test", + "use_tools": true, + "max_tool_rounds": 10 + }); + let parsed = RoutineAction::from_db("lightweight", json).expect("parse"); + match parsed { + RoutineAction::Lightweight { + max_tool_rounds, .. + } => { + assert_eq!(max_tool_rounds, 10, "normal value should pass through"); + } + _ => panic!("expected Lightweight"), + } + } } diff --git a/src/agent/routine_engine.rs b/src/agent/routine_engine.rs index b10021ef..b4aa5e0c 100644 --- a/src/agent/routine_engine.rs +++ b/src/agent/routine_engine.rs @@ -459,7 +459,20 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) prompt, context_paths, max_tokens, - } => execute_lightweight(&ctx, &routine, prompt, context_paths, *max_tokens).await, + use_tools, + max_tool_rounds, + } => { + execute_lightweight( + &ctx, + &routine, + prompt, + context_paths, + *max_tokens, + *use_tools, + *max_tool_rounds, + ) + .await + } RoutineAction::FullJob { title, description, @@ -670,6 +683,8 @@ async fn execute_lightweight( prompt: &str, context_paths: &[String], max_tokens: u32, + use_tools: bool, + max_tool_rounds: u32, ) -> Result<(RunStatus, Option, Option), RoutineError> { // Load context from workspace let mut context_parts = Vec::new(); @@ -732,14 +747,15 @@ async fn execute_lightweight( Err(_) => max_tokens, }; - // If tools are enabled, use the tool execution loop; otherwise, single LLM call - if ctx.config.lightweight_tools_enabled { + // If tools are enabled (both globally and per-routine), use the tool execution loop + if use_tools && ctx.config.lightweight_tools_enabled { execute_lightweight_with_tools( ctx, routine, &system_prompt, &full_prompt, effective_max_tokens, + max_tool_rounds, ) .await } else { @@ -783,24 +799,12 @@ async fn execute_lightweight_no_tools( reason: e.to_string(), })?; - let content = response.content.trim(); - let tokens_used = Some((response.input_tokens + response.output_tokens) as i32); - - // Empty content guard - if content.is_empty() { - return if response.finish_reason == FinishReason::Length { - Err(RoutineError::TruncatedResponse) - } else { - Err(RoutineError::EmptyResponse) - }; - } - - // Check for the "nothing to do" sentinel - if content == "ROUTINE_OK" || content.contains("ROUTINE_OK") { - return Ok((RunStatus::Ok, None, tokens_used)); - } - - Ok((RunStatus::Attention, Some(content.to_string()), tokens_used)) + handle_text_response( + &response.content, + response.finish_reason, + response.input_tokens, + response.output_tokens, + ) } /// Handle a text-only LLM response in lightweight routine execution. @@ -850,6 +854,7 @@ async fn execute_lightweight_with_tools( system_prompt: &str, full_prompt: &str, effective_max_tokens: u32, + max_tool_rounds: u32, ) -> Result<(RunStatus, Option, Option), RoutineError> { let mut messages = if system_prompt.is_empty() { vec![ChatMessage::user(full_prompt)] @@ -860,7 +865,9 @@ async fn execute_lightweight_with_tools( ] }; - let max_iterations = ctx.config.lightweight_max_iterations.min(5); + let max_iterations = max_tool_rounds + .min(ctx.config.lightweight_max_iterations) + .min(5); let mut iteration = 0; let mut total_input_tokens = 0; let mut total_output_tokens = 0; @@ -906,7 +913,10 @@ async fn execute_lightweight_with_tools( ); } else { // Tool-enabled iteration - let tool_defs = ctx.tools.tool_definitions().await; + let tool_defs = ctx + .tools + .tool_definitions_excluding(ROUTINE_TOOL_DENYLIST) + .await; let request = ToolCompletionRequest::new(messages.clone(), tool_defs) .with_max_tokens(effective_max_tokens) @@ -972,12 +982,33 @@ async fn execute_lightweight_with_tools( } } +/// Tools that must never be callable from lightweight routines. +/// +/// These tools pose autonomy-escalation risks: a routine could self-replicate, +/// modify its own triggers/prompts, delete other routines, or restart the agent. +const ROUTINE_TOOL_DENYLIST: &[&str] = &[ + "routine_create", + "routine_update", + "routine_delete", + "routine_fire", + "restart", +]; + /// Execute a single tool for a lightweight routine. async fn execute_routine_tool( ctx: &EngineContext, job_ctx: &JobContext, tc: &ToolCall, ) -> Result> { + // Block tools that pose autonomy-escalation risks + if ROUTINE_TOOL_DENYLIST.contains(&tc.name.as_str()) { + return Err(format!( + "Tool '{}' is not available in lightweight routines", + tc.name + ) + .into()); + } + // Check if tool exists let tool = ctx .tools @@ -1119,9 +1150,11 @@ pub fn spawn_cron_ticker( interval: Duration, ) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { + // Run one check immediately so routines due at startup don't wait + // an extra full polling interval. + engine.check_cron_triggers().await; + let mut ticker = tokio::time::interval(interval); - // Skip immediate first tick - ticker.tick().await; loop { ticker.tick().await; @@ -1283,6 +1316,36 @@ mod tests { } } + #[test] + fn test_routine_tool_denylist_blocks_self_management_tools() { + let denylisted = vec![ + "routine_create", + "routine_update", + "routine_delete", + "routine_fire", + "restart", + ]; + for tool in &denylisted { + assert!( + super::ROUTINE_TOOL_DENYLIST.contains(tool), + "Tool '{}' should be in ROUTINE_TOOL_DENYLIST", + tool + ); + } + } + + #[test] + fn test_routine_tool_denylist_allows_safe_tools() { + let allowed = vec!["echo", "time", "json", "http", "memory_search", "shell"]; + for tool in &allowed { + assert!( + !super::ROUTINE_TOOL_DENYLIST.contains(tool), + "Tool '{}' should NOT be in ROUTINE_TOOL_DENYLIST", + tool + ); + } + } + #[test] fn test_empty_response_handling() { // Simulate the empty content guard logic @@ -1297,4 +1360,11 @@ mod tests { assert_eq!(finish_reason_length, crate::llm::FinishReason::Length); assert_eq!(finish_reason_stop, crate::llm::FinishReason::Stop); } + + #[test] + fn test_truncate_adds_ellipsis_when_over_limit() { + let input = "abcdefghijk"; + let out = super::truncate(input, 5); + assert_eq!(out, "abcde..."); + } } diff --git a/src/app.rs b/src/app.rs index c53bebab..da77d3f3 100644 --- a/src/app.rs +++ b/src/app.rs @@ -9,6 +9,7 @@ use std::sync::Arc; +use crate::agent::SessionManager as AgentSessionManager; use crate::channels::web::log_layer::LogBroadcaster; use crate::config::Config; use crate::context::ContextManager; @@ -46,6 +47,8 @@ pub struct AppComponents { pub log_broadcaster: Arc, pub context_manager: Arc, pub hooks: Arc, + /// Shared thread/session manager used by the standard agent runtime. + pub agent_session_manager: Arc, pub skill_registry: Option>>, pub skill_catalog: Option>, pub cost_guard: Arc, @@ -287,6 +290,7 @@ impl AppBuilder { Arc::new(ToolRegistry::new()) }; tools.register_builtin_tools(); + tools.register_tool_info(); if let Some(ref ss) = self.secrets_store { tools.register_secrets_tools(Arc::clone(ss)); @@ -300,7 +304,8 @@ impl AppBuilder { // Register memory tools if database is available let workspace = if let Some(ref db) = self.db { - let mut ws = Workspace::new_with_db("default", db.clone()); + let mut ws = Workspace::new_with_db("default", db.clone()) + .with_search_config(&self.config.search); if let Some(ref emb) = embeddings { ws = ws.with_embeddings(emb.clone()); } @@ -689,6 +694,8 @@ impl AppBuilder { // Create hook registry early so runtime extension activation can register hooks. let hooks = Arc::new(HookRegistry::new()); + let agent_session_manager = + Arc::new(AgentSessionManager::new().with_hooks(Arc::clone(&hooks))); let ( mcp_session_manager, @@ -795,6 +802,7 @@ impl AppBuilder { log_broadcaster: self.log_broadcaster, context_manager, hooks, + agent_session_manager, skill_registry, skill_catalog, cost_guard, @@ -805,3 +813,69 @@ impl AppBuilder { }) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use async_trait::async_trait; + use tokio::sync::mpsc; + + use crate::agent::SessionManager as AgentSessionManager; + use crate::hooks::{ + Hook, HookContext, HookError, HookEvent, HookOutcome, HookPoint, HookRegistry, + }; + + struct SessionStartHook { + tx: mpsc::UnboundedSender<(String, String)>, + } + + #[async_trait] + impl Hook for SessionStartHook { + fn name(&self) -> &str { + "session-start-test" + } + + fn hook_points(&self) -> &[HookPoint] { + &[HookPoint::OnSessionStart] + } + + async fn execute( + &self, + event: &HookEvent, + _ctx: &HookContext, + ) -> Result { + if let HookEvent::SessionStart { + user_id, + session_id, + } = event + { + self.tx + .send((user_id.clone(), session_id.clone())) + .expect("test channel receiver should be alive"); + } else { + panic!("SessionStartHook received an unexpected event: {event:?}"); + } + Ok(HookOutcome::ok()) + } + } + + #[tokio::test] + async fn agent_session_manager_runs_session_start_hooks() { + let hooks = Arc::new(HookRegistry::new()); + let (tx, mut rx) = mpsc::unbounded_channel(); + hooks.register(Arc::new(SessionStartHook { tx })).await; + + let manager = AgentSessionManager::new().with_hooks(Arc::clone(&hooks)); + manager.get_or_create_session("user-123").await; + + let (user_id, session_id) = + tokio::time::timeout(std::time::Duration::from_secs(1), rx.recv()) + .await + .expect("session start hook should fire") + .expect("session start payload should be present"); + + assert_eq!(user_id, "user-123"); + assert!(!session_id.is_empty()); + } +} diff --git a/src/channels/http.rs b/src/channels/http.rs index 42fc54f8..15468c6a 100644 --- a/src/channels/http.rs +++ b/src/channels/http.rs @@ -269,95 +269,105 @@ async fn webhook_handler( let mut fallback_req = None; { let webhook_secret = state.webhook_secret.read().await; - if let Some(expected_secret) = webhook_secret.as_ref() { - let expected_secret = expected_secret.expose_secret(); + let Some(expected_secret) = webhook_secret.as_ref() else { + return ( + StatusCode::UNAUTHORIZED, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some( + "Webhook authentication required: HTTP webhook secret is not configured." + .to_string(), + ), + }), + ) + .into_response(); + }; + let expected_secret = expected_secret.expose_secret(); - match headers.get("x-ironclaw-signature") { - Some(raw_signature) => match raw_signature.to_str() { - Ok(signature) => { - if !verify_hmac_signature(expected_secret, &body, signature) { - return ( - StatusCode::UNAUTHORIZED, - Json(WebhookResponse { - message_id: Uuid::nil(), - status: "error".to_string(), - response: Some("Invalid webhook signature".to_string()), - }), - ) - .into_response(); - } + match headers.get("x-ironclaw-signature") { + Some(raw_signature) => match raw_signature.to_str() { + Ok(signature) => { + if !verify_hmac_signature(expected_secret, &body, signature) { + return ( + StatusCode::UNAUTHORIZED, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some("Invalid webhook signature".to_string()), + }), + ) + .into_response(); } + } + Err(_) => { + return ( + StatusCode::UNAUTHORIZED, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some("Invalid signature header encoding".to_string()), + }), + ) + .into_response(); + } + }, + None => { + let req: WebhookRequest = match serde_json::from_slice(&body) { + Ok(req) => req, Err(_) => { return ( StatusCode::UNAUTHORIZED, Json(WebhookResponse { message_id: Uuid::nil(), status: "error".to_string(), - response: Some("Invalid signature header encoding".to_string()), + response: Some( + "Webhook authentication required. Provide X-IronClaw-Signature header \ + (preferred) or 'secret' field in body (deprecated)." + .to_string(), + ), }), ) .into_response(); } - }, - None => { - let req: WebhookRequest = match serde_json::from_slice(&body) { - Ok(req) => req, - Err(_) => { - return ( - StatusCode::UNAUTHORIZED, - Json(WebhookResponse { - message_id: Uuid::nil(), - status: "error".to_string(), - response: Some( - "Webhook authentication required. Provide X-IronClaw-Signature header \ - (preferred) or 'secret' field in body (deprecated)." - .to_string(), - ), - }), - ) - .into_response(); - } - }; + }; - match &req.secret { - Some(provided) - if bool::from( - provided.as_bytes().ct_eq(expected_secret.as_bytes()), - ) => - { - tracing::warn!( - "Webhook authenticated via deprecated 'secret' field in request body. \ - Migrate to X-IronClaw-Signature header (HMAC-SHA256). \ - Body secret support will be removed in a future release." - ); - fallback_req = Some(req); - } - Some(_) => { - return ( - StatusCode::UNAUTHORIZED, - Json(WebhookResponse { - message_id: Uuid::nil(), - status: "error".to_string(), - response: Some("Invalid webhook secret".to_string()), - }), - ) - .into_response(); - } - None => { - return ( - StatusCode::UNAUTHORIZED, - Json(WebhookResponse { - message_id: Uuid::nil(), - status: "error".to_string(), - response: Some( - "Webhook authentication required. Provide X-IronClaw-Signature header \ - (preferred) or 'secret' field in body (deprecated)." - .to_string(), - ), - }), - ) - .into_response(); - } + match &req.secret { + Some(provided) + if bool::from(provided.as_bytes().ct_eq(expected_secret.as_bytes())) => + { + tracing::warn!( + "Webhook authenticated via deprecated 'secret' field in request body. \ + Migrate to X-IronClaw-Signature header (HMAC-SHA256). \ + Body secret support will be removed in a future release." + ); + fallback_req = Some(req); + } + Some(_) => { + return ( + StatusCode::UNAUTHORIZED, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some("Invalid webhook secret".to_string()), + }), + ) + .into_response(); + } + None => { + return ( + StatusCode::UNAUTHORIZED, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some( + "Webhook authentication required. Provide X-IronClaw-Signature header \ + (preferred) or 'secret' field in body (deprecated)." + .to_string(), + ), + }), + ) + .into_response(); } } } @@ -1052,6 +1062,32 @@ mod tests { ); } + #[tokio::test] + async fn webhook_rejects_requests_after_secret_is_cleared() { + let secret = "test-secret-123"; + let channel = test_channel(Some(secret)); + let _stream = channel.start().await.unwrap(); + let app = channel.routes(); + + channel.update_secret(None).await; + + let body = serde_json::json!({ + "content": "hello" + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + let signature = compute_signature(secret, &body_bytes); + let req = Request::builder() + .method("POST") + .uri("/webhook") + .header("content-type", "application/json") + .header("x-ironclaw-signature", signature) + .body(Body::from(body_bytes)) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + } + #[tokio::test] async fn test_concurrent_requests_during_secret_update() { use std::sync::Arc as StdArc; diff --git a/src/channels/relay/channel.rs b/src/channels/relay/channel.rs index cb64e882..d6aa90cc 100644 --- a/src/channels/relay/channel.rs +++ b/src/channels/relay/channel.rs @@ -408,12 +408,120 @@ impl Channel for RelayChannel { Ok(()) } - /// Status updates are not forwarded to messaging providers to avoid noise. async fn send_status( &self, - _status: StatusUpdate, - _metadata: &serde_json::Value, + status: StatusUpdate, + metadata: &serde_json::Value, ) -> Result<(), ChannelError> { + // Only handle ApprovalNeeded — all other variants are no-ops + let StatusUpdate::ApprovalNeeded { + request_id, + tool_name, + description, + parameters, + } = status + else { + return Ok(()); + }; + + // Only send buttons in DMs (dispatcher gates upstream, but guard here too) + let event_type = metadata + .get("event_type") + .and_then(|v| v.as_str()) + .unwrap_or(""); + if event_type != "direct_message" { + tracing::warn!( + tool = %tool_name, + event_type, + "Approval requested in non-DM, skipping buttons" + ); + return Ok(()); + } + + // Extract required metadata — error if missing + let channel_id = metadata + .get("channel_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ChannelError::SendFailed { + name: self.name().to_string(), + reason: "Missing channel_id for approval buttons".into(), + })?; + let sender_id = metadata + .get("sender_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ChannelError::SendFailed { + name: self.name().to_string(), + reason: "Missing sender_id for approval buttons".into(), + })?; + let thread_id = metadata.get("thread_id").and_then(|v| v.as_str()); + let team_id = metadata + .get("team_id") + .and_then(|v| v.as_str()) + .unwrap_or(&self.team_id); + + // Button value payload (Slack limits button values to 2000 chars; + // safe with typical UUIDs but documented here as a constraint) + let value_payload = serde_json::json!({ + "instance_id": self.instance_id, + "team_id": team_id, + "channel_id": channel_id, + "thread_ts": thread_id, + "request_id": request_id, + "sender_id": sender_id, + }); + let value_str = value_payload.to_string(); + + // Parameters are already redacted via redact_params() in dispatcher.rs + let params_display = + serde_json::to_string_pretty(¶meters).unwrap_or_else(|_| parameters.to_string()); + + let blocks = serde_json::json!([ + { + "type": "section", + "text": { + "type": "mrkdwn", + "text": format!( + "*Tool approval required*\n`{tool_name}`: {description}\n```{params_display}```" + ) + } + }, + { + "type": "actions", + "elements": [ + { + "type": "button", + "text": { "type": "plain_text", "text": "Approve" }, + "style": "primary", + "action_id": "approve_tool", + "value": value_str, + }, + { + "type": "button", + "text": { "type": "plain_text", "text": "Deny" }, + "style": "danger", + "action_id": "deny_tool", + "value": value_str, + } + ] + } + ]); + + let mut body = serde_json::json!({ + "channel": channel_id, + "text": format!("Tool approval required: {tool_name} - {description}"), + "blocks": blocks, + }); + if let Some(tid) = thread_id { + body["thread_ts"] = serde_json::Value::String(tid.to_string()); + } + + self.proxy_send(team_id, "chat.postMessage", body) + .await + .map_err(|e| ChannelError::SendFailed { + name: self.name().to_string(), + reason: e.to_string(), + })?; + Ok(()) } @@ -639,4 +747,118 @@ mod tests { // The reconnect loop now skips team validation when team_id is empty, // so the channel remains alive. } + + #[tokio::test] + async fn test_send_status_non_approval_is_noop() { + let channel = RelayChannel::new( + test_client(), + "token".into(), + "T123".into(), + "inst1".into(), + "user1".into(), + ); + let metadata = serde_json::json!({}); + let result = channel + .send_status( + StatusUpdate::ToolStarted { + name: "echo".into(), + }, + &metadata, + ) + .await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_send_status_approval_non_dm_skips() { + let channel = RelayChannel::new( + test_client(), + "token".into(), + "T123".into(), + "inst1".into(), + "user1".into(), + ); + let metadata = serde_json::json!({ + "event_type": "message", + "channel_id": "C456", + "sender_id": "U789", + }); + let result = channel + .send_status( + StatusUpdate::ApprovalNeeded { + request_id: "req1".into(), + tool_name: "shell".into(), + description: "run command".into(), + parameters: serde_json::json!({}), + }, + &metadata, + ) + .await; + // Non-DM approval requests are silently skipped (no HTTP call) + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_send_status_approval_dm_missing_channel_id_errors() { + let channel = RelayChannel::new( + test_client(), + "token".into(), + "T123".into(), + "inst1".into(), + "user1".into(), + ); + let metadata = serde_json::json!({ + "event_type": "direct_message", + "sender_id": "U789", + }); + let result = channel + .send_status( + StatusUpdate::ApprovalNeeded { + request_id: "req1".into(), + tool_name: "shell".into(), + description: "run command".into(), + parameters: serde_json::json!({}), + }, + &metadata, + ) + .await; + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("channel_id"), + "expected channel_id error, got: {err}" + ); + } + + #[tokio::test] + async fn test_send_status_approval_dm_missing_sender_id_errors() { + let channel = RelayChannel::new( + test_client(), + "token".into(), + "T123".into(), + "inst1".into(), + "user1".into(), + ); + let metadata = serde_json::json!({ + "event_type": "direct_message", + "channel_id": "C456", + }); + let result = channel + .send_status( + StatusUpdate::ApprovalNeeded { + request_id: "req1".into(), + tool_name: "shell".into(), + description: "run command".into(), + parameters: serde_json::json!({}), + }, + &metadata, + ) + .await; + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("sender_id"), + "expected sender_id error, got: {err}" + ); + } } diff --git a/src/channels/web/handlers/routines.rs b/src/channels/web/handlers/routines.rs index d8803efa..f49d7fe8 100644 --- a/src/channels/web/handlers/routines.rs +++ b/src/channels/web/handlers/routines.rs @@ -10,6 +10,7 @@ use axum::{ use serde::Deserialize; use uuid::Uuid; +use crate::agent::routine::{Trigger, next_cron_fire}; use crate::channels::web::server::GatewayState; use crate::channels::web::types::*; use crate::error::RoutineError; @@ -182,12 +183,21 @@ pub async fn routines_toggle_handler( .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? .ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?; + let was_enabled = routine.enabled; // If a specific value was provided, use it; otherwise toggle. routine.enabled = match body { Some(Json(req)) => req.enabled.unwrap_or(!routine.enabled), None => !routine.enabled, }; + if routine.enabled + && !was_enabled + && let Trigger::Cron { schedule, timezone } = &routine.trigger + { + routine.next_fire_at = next_cron_fire(schedule, timezone.as_deref()) + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + } + store .update_routine(&routine) .await diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index 904971fc..48ef452c 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -26,6 +26,7 @@ use tower_http::set_header::SetResponseHeaderLayer; use uuid::Uuid; use crate::agent::SessionManager; +use crate::agent::routine::{Trigger, next_cron_fire}; use crate::bootstrap::ironclaw_base_dir; use crate::channels::IncomingMessage; use crate::channels::relay::DEFAULT_RELAY_NAME; @@ -572,6 +573,14 @@ async fn oauth_callback_handler( extension = %flow.extension_name, "OAuth flow expired" ); + // Notify UI so auth card can show error instead of staying stuck + if let Some(ref sender) = flow.sse_sender { + let _ = sender.send(SseEvent::AuthCompleted { + extension_name: flow.extension_name.clone(), + success: false, + message: "OAuth flow expired. Please try again.".to_string(), + }); + } return oauth_error_page(&flow.display_name); } @@ -2416,12 +2425,21 @@ async fn routines_toggle_handler( .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? .ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?; + let was_enabled = routine.enabled; // If a specific value was provided, use it; otherwise toggle. routine.enabled = match body { Some(Json(req)) => req.enabled.unwrap_or(!routine.enabled), None => !routine.enabled, }; + if routine.enabled + && !was_enabled + && let Trigger::Cron { schedule, timezone } = &routine.trigger + { + routine.next_fire_at = next_cron_fire(schedule, timezone.as_deref()) + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + } + store .update_routine(&routine) .await @@ -2696,6 +2714,7 @@ struct GatewayStatusResponse { #[cfg(test)] mod tests { use super::*; + use crate::cli::oauth_defaults; use crate::testing::credentials::TEST_GATEWAY_CRYPTO_KEY; #[test] @@ -2813,6 +2832,11 @@ mod tests { .with_state(state) } + fn expired_flow_created_at() -> Option { + std::time::Instant::now() + .checked_sub(oauth_defaults::OAUTH_FLOW_EXPIRY + std::time::Duration::from_secs(1)) + } + #[tokio::test] async fn test_csp_header_present_on_responses() { use std::net::SocketAddr; @@ -2919,29 +2943,14 @@ mod tests { use tower::ServiceExt; // Build an ExtensionManager so the handler can look up flows - let secrets = Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new( - crate::secrets::SecretsCrypto::new(secrecy::SecretString::from( - TEST_GATEWAY_CRYPTO_KEY.to_string(), - )) - .expect("crypto"), - ))); - let tool_registry = Arc::new(ToolRegistry::new()); - let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new()); - - let ext_mgr = Arc::new(ExtensionManager::new( - mcp_sm, - Arc::new(crate::tools::mcp::process::McpProcessManager::new()), - secrets, - tool_registry, - None, - None, - std::path::PathBuf::from("/tmp/wasm_tools"), - std::path::PathBuf::from("/tmp/wasm_channels"), - None, - "test".to_string(), - None, - vec![], - )); + let secrets: Arc = + Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new( + crate::secrets::SecretsCrypto::new(secrecy::SecretString::from( + TEST_GATEWAY_CRYPTO_KEY.to_string(), + )) + .expect("crypto"), + ))); + let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets); let state = test_gateway_state(Some(ext_mgr)); let app = test_oauth_router(state); @@ -2975,25 +2984,13 @@ mod tests { )) .expect("crypto"), ))); - let tool_registry = Arc::new(ToolRegistry::new()); - let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new()); + let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone()); + let Some(created_at) = expired_flow_created_at() else { + eprintln!("Skipping expired OAuth flow test: monotonic uptime below expiry window"); + return; + }; - let ext_mgr = Arc::new(ExtensionManager::new( - mcp_sm, - Arc::new(crate::tools::mcp::process::McpProcessManager::new()), - secrets.clone(), - tool_registry, - None, - None, - std::path::PathBuf::from("/tmp/wasm_tools"), - std::path::PathBuf::from("/tmp/wasm_channels"), - None, - "test".to_string(), - None, - vec![], - )); - - // Insert an expired flow (created 10 minutes ago) + // Insert an expired flow. let flow = crate::cli::oauth_defaults::PendingOAuthFlow { extension_name: "test_tool".to_string(), display_name: "Test Tool".to_string(), @@ -3013,9 +3010,7 @@ mod tests { gateway_token: None, resource: None, client_id_secret_name: None, - created_at: std::time::Instant::now() - .checked_sub(std::time::Duration::from_secs(600)) - .expect("System uptime is too low to run expired flow test"), + created_at, }; ext_mgr @@ -3045,6 +3040,80 @@ mod tests { assert!(html.contains("Authorization Failed")); } + #[tokio::test] + async fn test_oauth_callback_expired_flow_broadcasts_auth_completed_failure() { + use axum::body::Body; + use tower::ServiceExt; + + let secrets: Arc = + Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new( + crate::secrets::SecretsCrypto::new(secrecy::SecretString::from( + TEST_GATEWAY_CRYPTO_KEY.to_string(), + )) + .expect("crypto"), + ))); + let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone()); + + let (sender, mut receiver) = tokio::sync::broadcast::channel(4); + let Some(created_at) = expired_flow_created_at() else { + eprintln!("Skipping expired OAuth flow SSE test: monotonic uptime below expiry window"); + return; + }; + let flow = crate::cli::oauth_defaults::PendingOAuthFlow { + extension_name: "test_tool".to_string(), + display_name: "Test Tool".to_string(), + token_url: "https://example.com/token".to_string(), + client_id: "client123".to_string(), + client_secret: None, + redirect_uri: "https://example.com/oauth/callback".to_string(), + code_verifier: None, + access_token_field: "access_token".to_string(), + secret_name: "test_token".to_string(), + provider: None, + validation_endpoint: None, + scopes: vec![], + user_id: "test".to_string(), + secrets, + sse_sender: Some(sender), + gateway_token: None, + resource: None, + client_id_secret_name: None, + created_at, + }; + + ext_mgr + .pending_oauth_flows() + .write() + .await + .insert("expired_state".to_string(), flow); + + let state = test_gateway_state(Some(ext_mgr)); + let app = test_oauth_router(state); + + let req = axum::http::Request::builder() + .uri("/oauth/callback?code=test_code&state=expired_state") + .body(Body::empty()) + .expect("request"); + + let resp = ServiceExt::>::oneshot(app, req) + .await + .expect("response"); + assert_eq!(resp.status(), StatusCode::OK); + + match receiver.recv().await.expect("auth_completed event") { + crate::channels::web::types::SseEvent::AuthCompleted { + extension_name, + success, + message, + } => { + assert_eq!(extension_name, "test_tool"); + assert!(!success, "expired OAuth flow should broadcast failure"); + assert_eq!(message, "OAuth flow expired. Please try again."); + } + event => panic!("expected AuthCompleted event, got {event:?}"), + } + } + #[tokio::test] async fn test_oauth_callback_no_extension_manager() { use axum::body::Body; @@ -3083,28 +3152,16 @@ mod tests { )) .expect("crypto"), ))); - let tool_registry = Arc::new(ToolRegistry::new()); - let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new()); - - let ext_mgr = Arc::new(ExtensionManager::new( - mcp_sm, - Arc::new(crate::tools::mcp::process::McpProcessManager::new()), - secrets.clone(), - tool_registry, - None, - None, - std::path::PathBuf::from("/tmp/wasm_tools"), - std::path::PathBuf::from("/tmp/wasm_channels"), - None, - "test".to_string(), - None, - vec![], - )); + let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone()); // Insert a flow keyed by raw nonce "test_nonce" (without instance prefix). // Use an expired flow so the handler exits before attempting a real HTTP // token exchange — we only need to verify that the instance prefix was // stripped and the flow was found by the raw nonce. + let Some(created_at) = expired_flow_created_at() else { + eprintln!("Skipping OAuth state-prefix test: monotonic uptime below expiry window"); + return; + }; let flow = crate::cli::oauth_defaults::PendingOAuthFlow { extension_name: "test_tool".to_string(), display_name: "Test Tool".to_string(), @@ -3125,9 +3182,7 @@ mod tests { resource: None, client_id_secret_name: None, // Expired — handler will reject after lookup (no network I/O) - created_at: std::time::Instant::now() - .checked_sub(std::time::Duration::from_secs(600)) - .expect("System uptime is too low to run expired flow test"), + created_at, }; ext_mgr @@ -3198,24 +3253,27 @@ mod tests { fn test_ext_mgr( secrets: Arc, - ) -> Arc { + ) -> (Arc, tempfile::TempDir, tempfile::TempDir) { let tool_registry = Arc::new(ToolRegistry::new()); let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new()); let mcp_pm = Arc::new(crate::tools::mcp::process::McpProcessManager::new()); - Arc::new(ExtensionManager::new( + let wasm_tools_dir = tempfile::tempdir().expect("temp wasm tools dir"); + let wasm_channels_dir = tempfile::tempdir().expect("temp wasm channels dir"); + let ext_mgr = Arc::new(ExtensionManager::new( mcp_sm, mcp_pm, secrets, tool_registry, None, None, - std::path::PathBuf::from("/tmp/wasm_tools"), - std::path::PathBuf::from("/tmp/wasm_channels"), + wasm_tools_dir.path().to_path_buf(), + wasm_channels_dir.path().to_path_buf(), None, "test".to_string(), None, vec![], - )) + )); + (ext_mgr, wasm_tools_dir, wasm_channels_dir) } #[tokio::test] @@ -3224,7 +3282,7 @@ mod tests { use tower::ServiceExt; let secrets = test_secrets_store(); - let ext_mgr = test_ext_mgr(secrets); + let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets); let state = test_gateway_state(Some(ext_mgr)); let app = test_relay_oauth_router(state); @@ -3268,7 +3326,7 @@ mod tests { .await .expect("store nonce"); - let ext_mgr = test_ext_mgr(secrets); + let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets); let state = test_gateway_state(Some(ext_mgr)); let app = test_relay_oauth_router(state); @@ -3313,7 +3371,7 @@ mod tests { .await .expect("store nonce"); - let ext_mgr = test_ext_mgr(secrets.clone()); + let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone()); let state = test_gateway_state(Some(ext_mgr)); let app = test_relay_oauth_router(state); diff --git a/src/channels/web/static/app.js b/src/channels/web/static/app.js index f5812030..081ae60d 100644 --- a/src/channels/web/static/app.js +++ b/src/channels/web/static/app.js @@ -342,31 +342,27 @@ function connectSSE() { eventSource.addEventListener('approval_needed', (e) => { const data = JSON.parse(e.data); - if (!isCurrentThread(data.thread_id)) return; - showApproval(data); + const hasThread = !!data.thread_id; + const forCurrentThread = !hasThread || isCurrentThread(data.thread_id); + + if (forCurrentThread) { + showApproval(data); + } else { + // Keep thread list fresh when approval is requested in a background thread. + unreadThreads.set(data.thread_id, (unreadThreads.get(data.thread_id) || 0) + 1); + debouncedLoadThreads(); + } + + // Extension setup flows can surface approvals while user is on Extensions tab. + if (currentTab === 'extensions') loadExtensions(); }); eventSource.addEventListener('auth_required', (e) => { - const data = JSON.parse(e.data); - if (data.auth_url) { - // OAuth flow: show the auth card with an OAuth button + optional token paste field. - showAuthCard(data); - } else { - // Setup flow: fetch the extension's credential schema and show the multi-field - // configure modal (the same UI used by the Extensions tab "Setup" button). - showConfigureModal(data.extension_name); - } + handleAuthRequired(JSON.parse(e.data)); }); eventSource.addEventListener('auth_completed', (e) => { - const data = JSON.parse(e.data); - // Dismiss whichever UI path was active: auth card (OAuth) or configure modal (setup). - removeAuthCard(data.extension_name); - closeConfigureModal(); - showToast(data.message, data.success ? 'success' : 'error'); - // Refresh extensions list so status indicators update - if (currentTab === 'extensions') loadExtensions(); - enableChatInput(); + handleAuthCompleted(JSON.parse(e.data)); }); eventSource.addEventListener('extension_status', (e) => { @@ -702,16 +698,25 @@ function copyCodeBlock(btn) { }); } +function copyMessage(btn) { + const message = btn.closest('.message'); + if (!message) return; + const text = message.getAttribute('data-copy-text') + || message.getAttribute('data-raw') + || message.textContent + || ''; + navigator.clipboard.writeText(text).then(() => { + btn.textContent = 'Copied'; + setTimeout(() => { btn.textContent = 'Copy'; }, 1200); + }).catch(() => { + btn.textContent = 'Failed'; + setTimeout(() => { btn.textContent = 'Copy'; }, 1200); + }); +} + function addMessage(role, content) { const container = document.getElementById('chat-messages'); - const div = document.createElement('div'); - div.className = 'message ' + role; - if (role === 'user') { - div.textContent = content; - } else { - div.setAttribute('data-raw', content); - div.innerHTML = renderMarkdown(content); - } + const div = createMessageElement(role, content); container.appendChild(div); container.scrollTop = container.scrollHeight; } @@ -723,7 +728,11 @@ function appendToLastAssistant(chunk) { const last = messages[messages.length - 1]; const raw = (last.getAttribute('data-raw') || '') + chunk; last.setAttribute('data-raw', raw); - last.innerHTML = renderMarkdown(raw); + last.setAttribute('data-copy-text', raw); + const content = last.querySelector('.message-content'); + if (content) { + content.innerHTML = renderMarkdown(raw); + } container.scrollTop = container.scrollHeight; } else { addMessage('assistant', chunk); @@ -977,7 +986,26 @@ function finalizeActivityGroup() { _activeToolCards = {}; } +function humanizeToolName(rawName) { + if (!rawName) return ''; + return String(rawName) + .replace(/[_-]+/g, ' ') + .replace(/([a-z0-9])([A-Z])/g, '$1 $2') + .replace(/^tool([a-zA-Z])/, 'tool $1') + .replace(/\s+/g, ' ') + .trim(); +} + +function shouldShowChannelConnectedMessage(extensionName, success) { + if (!success || !extensionName) return false; + return String(extensionName).toLowerCase().includes('telegram'); +} + function showApproval(data) { + // Avoid duplicate cards on reconnect/history refresh. + const existing = document.querySelector('.approval-card[data-request-id="' + CSS.escape(data.request_id) + '"]'); + if (existing) return; + const container = document.getElementById('chat-messages'); const card = document.createElement('div'); card.className = 'approval-card'; @@ -990,7 +1018,7 @@ function showApproval(data) { const toolName = document.createElement('div'); toolName.className = 'approval-tool-name'; - toolName.textContent = data.tool_name; + toolName.textContent = humanizeToolName(data.tool_name); card.appendChild(toolName); if (data.description) { @@ -1093,13 +1121,71 @@ function showJobCard(data) { // --- Auth card --- +function handleAuthRequired(data) { + if (data.auth_url) { + // OAuth flow: show the global auth prompt with an OAuth button + optional token paste field. + showAuthCard(data); + } else { + // Setup flow: fetch the extension's credential schema and show the multi-field + // configure modal (the same UI used by the Extensions tab "Setup" button). + showConfigureModal(data.extension_name); + } +} + +function handleAuthCompleted(data) { + // Dismiss only the matching extension's UI so unrelated setup work is not interrupted. + removeAuthCard(data.extension_name); + closeConfigureModal(data.extension_name); + showToast(data.message, data.success ? 'success' : 'error'); + if (shouldShowChannelConnectedMessage(data.extension_name, data.success)) { + addMessage('system', 'Telegram is now connected. You can message me there and I can send you notifications.'); + } + if (currentTab === 'extensions') loadExtensions(); + enableChatInput(); +} + +function queryByDataAttribute(selector, attributeName, attributeValue) { + if (typeof attributeValue !== 'string') return document.querySelector(selector); + + if (window.CSS && typeof window.CSS.escape === 'function') { + return document.querySelector( + selector + '[' + attributeName + '="' + window.CSS.escape(attributeValue) + '"]' + ); + } + + const candidates = document.querySelectorAll(selector); + for (const candidate of candidates) { + if (candidate.getAttribute(attributeName) === attributeValue) return candidate; + } + return null; +} + +function getAuthOverlay(extensionName) { + return queryByDataAttribute('.auth-overlay', 'data-extension-name', extensionName); +} + +function getAuthCard(extensionName) { + return queryByDataAttribute('.auth-card', 'data-extension-name', extensionName); +} + +function getConfigureOverlay(extensionName) { + return queryByDataAttribute('.configure-overlay', 'data-extension-name', extensionName); +} + function showAuthCard(data) { - // Remove any existing card for this extension first - removeAuthCard(data.extension_name); + // Keep a single global auth prompt so the experience is consistent across tabs. + const existing = getAuthOverlay(); + if (existing) existing.remove(); + + const overlay = document.createElement('div'); + overlay.className = 'auth-overlay'; + overlay.setAttribute('data-extension-name', data.extension_name); + overlay.addEventListener('click', (e) => { + if (e.target === overlay) cancelAuth(data.extension_name); + }); - const container = document.getElementById('chat-messages'); const card = document.createElement('div'); - card.className = 'auth-card'; + card.className = 'auth-card auth-modal'; card.setAttribute('data-extension-name', data.extension_name); const header = document.createElement('div'); @@ -1178,21 +1264,30 @@ function showAuthCard(data) { actions.appendChild(cancelBtn); card.appendChild(actions); - container.appendChild(card); - container.scrollTop = container.scrollHeight; + overlay.appendChild(card); + document.body.appendChild(overlay); tokenInput.focus(); } function removeAuthCard(extensionName) { - const card = document.querySelector('.auth-card[data-extension-name="' + extensionName + '"]'); - if (card) card.remove(); + const overlay = getAuthOverlay(extensionName); + if (overlay) { + overlay.remove(); + return; + } + const card = getAuthCard(extensionName); + if (card) { + const parentOverlay = card.closest('.auth-overlay'); + if (parentOverlay) parentOverlay.remove(); + else card.remove(); + } } function submitAuthToken(extensionName, tokenValue) { if (!tokenValue || !tokenValue.trim()) return; // Disable submit button while in flight - const card = document.querySelector('.auth-card[data-extension-name="' + extensionName + '"]'); + const card = getAuthCard(extensionName); if (card) { const btns = card.querySelectorAll('button'); btns.forEach((b) => { b.disabled = true; }); @@ -1203,8 +1298,10 @@ function submitAuthToken(extensionName, tokenValue) { body: { extension_name: extensionName, token: tokenValue.trim() }, }).then((result) => { if (result.success) { + // Close immediately for responsiveness; the authoritative success UX + // (toast + extensions refresh) still comes from auth_completed SSE. removeAuthCard(extensionName); - addMessage('system', result.message); + enableChatInput(); } else { showAuthCardError(extensionName, result.message); } @@ -1223,7 +1320,7 @@ function cancelAuth(extensionName) { } function showAuthCardError(extensionName, message) { - const card = document.querySelector('.auth-card[data-extension-name="' + extensionName + '"]'); + const card = getAuthCard(extensionName); if (!card) return; // Re-enable buttons const btns = card.querySelectorAll('button'); @@ -1310,12 +1407,31 @@ function loadHistory(before) { function createMessageElement(role, content) { const div = document.createElement('div'); div.className = 'message ' + role; - if (role === 'user') { - div.textContent = content; + + if (role === 'assistant' || role === 'user') { + div.classList.add('has-copy'); + div.setAttribute('data-copy-text', content); + const copyBtn = document.createElement('button'); + copyBtn.className = 'message-copy-btn'; + copyBtn.type = 'button'; + copyBtn.setAttribute('aria-label', 'Copy message'); + copyBtn.textContent = 'Copy'; + copyBtn.addEventListener('click', (e) => { + e.stopPropagation(); + copyMessage(copyBtn); + }); + div.appendChild(copyBtn); + } + + const body = document.createElement('div'); + body.className = 'message-content'; + if (role === 'user' || role === 'system') { + body.textContent = content; } else { div.setAttribute('data-raw', content); - div.innerHTML = renderMarkdown(content); + body.innerHTML = renderMarkdown(content); } + div.appendChild(body); return div; } @@ -2134,6 +2250,10 @@ function renderAvailableExtensionCard(entry) { showToast(I18n.t('extensions.installedSuccess', {name: entry.display_name}), 'success'); // OAuth popup if auth started during install (builtin creds) if (res.auth_url) { + showAuthCard({ + extension_name: entry.name, + auth_url: res.auth_url, + }); showToast('Opening authentication for ' + entry.display_name, 'info'); openOAuthUrl(res.auth_url); } @@ -2399,6 +2519,10 @@ function activateExtension(name) { if (res.success) { // Even on success, the tool may need OAuth (e.g., WASM loaded but no token yet) if (res.auth_url) { + showAuthCard({ + extension_name: name, + auth_url: res.auth_url, + }); showToast('Opening authentication for ' + name, 'info'); openOAuthUrl(res.auth_url); } @@ -2407,6 +2531,10 @@ function activateExtension(name) { } if (res.auth_url) { + showAuthCard({ + extension_name: name, + auth_url: res.auth_url, + }); showToast('Opening authentication for ' + name, 'info'); openOAuthUrl(res.auth_url); } else if (res.awaiting_token) { @@ -2449,6 +2577,7 @@ function renderConfigureModal(name, secrets) { closeConfigureModal(); const overlay = document.createElement('div'); overlay.className = 'configure-overlay'; + overlay.setAttribute('data-extension-name', name); overlay.addEventListener('click', (e) => { if (e.target === overlay) closeConfigureModal(); }); @@ -2542,7 +2671,8 @@ function submitConfigureModal(name, fields) { } // Disable buttons to prevent double-submit - var btns = document.querySelectorAll('.configure-actions button'); + const overlay = getConfigureOverlay(name) || document.querySelector('.configure-overlay'); + var btns = overlay ? overlay.querySelectorAll('.configure-actions button') : []; btns.forEach(function(b) { b.disabled = true; }); apiFetch('/api/extensions/' + encodeURIComponent(name) + '/setup', { @@ -2553,8 +2683,10 @@ function submitConfigureModal(name, fields) { if (res.success) { closeConfigureModal(); if (res.auth_url) { - // OAuth flow started — open consent popup. The auth_completed SSE will - // not arrive immediately (it fires after OAuth callback), so show a toast now. + showAuthCard({ + extension_name: name, + auth_url: res.auth_url, + }); showToast('Opening OAuth authorization for ' + name, 'info'); openOAuthUrl(res.auth_url); loadExtensions(); @@ -2573,8 +2705,9 @@ function submitConfigureModal(name, fields) { }); } -function closeConfigureModal() { - const existing = document.querySelector('.configure-overlay'); +function closeConfigureModal(extensionName) { + if (typeof extensionName !== 'string') extensionName = null; + const existing = getConfigureOverlay(extensionName); if (existing) existing.remove(); } diff --git a/src/channels/web/static/style.css b/src/channels/web/static/style.css index a0985ce3..b6e1cbdf 100644 --- a/src/channels/web/static/style.css +++ b/src/channels/web/static/style.css @@ -666,6 +666,7 @@ body { font-size: 14px; line-height: 1.5; word-wrap: break-word; + position: relative; } .message.user { @@ -686,6 +687,58 @@ body { line-height: 1.6; } +.message.has-copy { + padding-right: 52px; +} + +.message-content { + min-width: 0; +} + +.message-copy-btn { + position: absolute; + top: 8px; + right: 8px; + z-index: 2; + border: 1px solid var(--border); + background: var(--bg-primary); + color: var(--text-secondary); + border-radius: 8px; + font-size: 11px; + padding: 2px 8px; + opacity: 0; + pointer-events: none; + transition: opacity 0.15s ease; +} + +.message.user:hover .message-copy-btn, +.message.assistant:hover .message-copy-btn, +.message.user:focus-within .message-copy-btn, +.message.assistant:focus-within .message-copy-btn { + opacity: 1; + pointer-events: auto; +} + +.message-copy-btn:focus-visible { + opacity: 1; + pointer-events: auto; + outline: 2px solid var(--accent); + outline-offset: 1px; +} + +.message-copy-btn:hover { + background: var(--bg-secondary); + color: var(--text-primary); +} + +@media (hover: none) { + .message.user .message-copy-btn, + .message.assistant .message-copy-btn { + opacity: 1; + pointer-events: auto; + } +} + .message.system { align-self: center; background: var(--bg-tertiary); @@ -1166,7 +1219,21 @@ body { color: var(--danger); } -/* Auth card (inline in chat) */ +/* Auth prompt */ +.auth-overlay { + position: fixed; + top: 0; + left: 0; + width: 100%; + height: 100%; + background: rgba(0, 0, 0, 0.6); + z-index: 1001; + display: flex; + align-items: center; + justify-content: center; + padding: 16px; +} + .auth-card { align-self: flex-start; max-width: 80%; @@ -1181,6 +1248,16 @@ body { transition: border-color 0.2s; } +.auth-overlay .auth-card { + width: 460px; + max-width: min(460px, 90vw); + margin: 0; + align-self: auto; + background: var(--bg); + border-color: rgba(52, 211, 153, 0.35); + box-shadow: 0 24px 48px rgba(0, 0, 0, 0.35); +} + .auth-card .auth-header { font-weight: 600; color: var(--accent); diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 7efab882..652cac01 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -7,6 +7,7 @@ //! - Managing WASM tools (`tool install`, `tool list`, `tool remove`) //! - Managing MCP servers (`mcp add`, `mcp auth`, `mcp list`, `mcp test`) //! - Querying workspace memory (`memory search`, `memory read`, `memory write`) +//! - Managing routines (`routines list`, `routines create`, `routines edit`, ...) //! - Managing OS service (`service install`, `service start`, `service stop`) //! - Listing configured channels (`channels list`) //! - Active health diagnostics (`doctor`) @@ -23,6 +24,7 @@ pub mod memory; pub mod oauth_defaults; mod pairing; mod registry; +mod routines; mod service; mod skills; pub mod status; @@ -39,6 +41,7 @@ pub use memory::MemoryCommand; pub use memory::run_memory_command_with_db; pub use pairing::{PairingCommand, run_pairing_command, run_pairing_command_with_store}; pub use registry::{RegistryCommand, run_registry_command}; +pub use routines::{RoutinesCommand, run_routines_command}; pub use service::{ServiceCommand, run_service_command}; pub use skills::{SkillsCommand, run_skills_command}; pub use status::run_status_command; @@ -147,6 +150,15 @@ pub enum Command { )] Channels(ChannelsCommand), + /// Manage routines (scheduled, event-driven, webhook, manual) + #[command( + subcommand, + alias = "cron", + about = "Manage routines", + long_about = "List, create, edit, enable/disable, delete, and view history of routines.\nExamples:\n ironclaw routines list\n ironclaw routines create --name daily-digest --schedule '0 0 9 * * *' --prompt 'Summarize today'" + )] + Routines(RoutinesCommand), + /// Manage MCP servers (hosted tool providers) #[command( subcommand, @@ -281,6 +293,23 @@ pub async fn init_secrets_store() Ok(crate::db::create_secrets_store(&config.database, crypto).await?) } +/// Run the Routines CLI subcommand. +pub async fn run_routines_cli( + routines_cmd: &RoutinesCommand, + config_path: Option<&std::path::Path>, +) -> anyhow::Result<()> { + let config = crate::config::Config::from_env_with_toml(config_path) + .await + .map_err(|e| anyhow::anyhow!("{e:#}"))?; + + let db: Arc = crate::db::connect_from_config(&config.database) + .await + .map_err(|e| anyhow::anyhow!("{e:#}"))?; + + let user_id = std::env::var("GATEWAY_USER_ID").unwrap_or_else(|_| "default".to_string()); + run_routines_command(routines_cmd.clone(), db, &user_id).await +} + /// Run the Memory CLI subcommand. pub async fn run_memory_command(mem_cmd: &MemoryCommand) -> anyhow::Result<()> { let config = crate::config::Config::from_env() diff --git a/src/cli/routines.rs b/src/cli/routines.rs new file mode 100644 index 00000000..852fc41f --- /dev/null +++ b/src/cli/routines.rs @@ -0,0 +1,732 @@ +//! `ironclaw routines` — manage scheduled routines from the CLI. +//! +//! Provides subcommands for listing, creating, editing, enabling/disabling, +//! deleting, and viewing run history of routines without starting the full agent. + +use std::sync::Arc; + +use chrono::{DateTime, Utc}; +use clap::Subcommand; +use uuid::Uuid; + +use crate::agent::routine::{ + NotifyConfig, Routine, RoutineAction, RoutineGuardrails, Trigger, next_cron_fire, +}; +use crate::db::Database; + +/// Routines subcommands. +#[derive(Subcommand, Debug, Clone)] +pub enum RoutinesCommand { + /// List routines + List { + /// Filter by trigger type (e.g. "cron", "webhook", "event") + #[arg(long)] + trigger: Option, + + /// Include disabled routines + #[arg(long)] + disabled: bool, + + /// Output as JSON (for scripting) + #[arg(long)] + json: bool, + }, + + /// Create a new cron routine + #[command(alias = "add")] + Create { + /// Routine name (must be unique per user) + #[arg(long)] + name: String, + + /// Cron schedule (6-field: "sec min hour day month weekday") + #[arg(long)] + schedule: String, + + /// Prompt for the LLM + #[arg(long)] + prompt: String, + + /// Optional description + #[arg(long, default_value = "")] + description: String, + + /// IANA timezone (e.g. "America/New_York") + #[arg(long)] + timezone: Option, + + /// Cooldown between fires in seconds + #[arg(long, default_value = "300")] + cooldown: u64, + + /// Notification channel + #[arg(long)] + notify_channel: Option, + }, + + /// Edit an existing routine + #[command(alias = "update")] + Edit { + /// Routine name + #[arg(long)] + name: String, + + /// New schedule + #[arg(long)] + schedule: Option, + + /// New prompt + #[arg(long)] + prompt: Option, + + /// New description + #[arg(long)] + description: Option, + + /// New timezone + #[arg(long)] + timezone: Option, + + /// New cooldown in seconds + #[arg(long)] + cooldown: Option, + }, + + /// Enable a routine + Enable { + /// Routine name + name: String, + }, + + /// Disable a routine + Disable { + /// Routine name + name: String, + }, + + /// Delete a routine + #[command(alias = "rm")] + Delete { + /// Routine name + name: String, + + /// Skip confirmation prompt + #[arg(short, long)] + yes: bool, + }, + + /// Show run history for a routine + #[command(alias = "runs")] + History { + /// Routine name + name: String, + + /// Maximum number of runs to show + #[arg(short, long, default_value = "10")] + limit: i64, + + /// Output as JSON (for scripting) + #[arg(long)] + json: bool, + }, +} + +/// Run a routines CLI command against the database. +pub async fn run_routines_command( + cmd: RoutinesCommand, + db: Arc, + user_id: &str, +) -> anyhow::Result<()> { + match cmd { + RoutinesCommand::List { + trigger, + disabled, + json, + } => list(&db, user_id, trigger.as_deref(), disabled, json).await, + RoutinesCommand::Create { + name, + schedule, + prompt, + description, + timezone, + cooldown, + notify_channel, + } => { + create( + &db, + user_id, + &name, + &schedule, + &prompt, + &description, + timezone.as_deref(), + cooldown, + notify_channel, + ) + .await + } + RoutinesCommand::Edit { + name, + schedule, + prompt, + description, + timezone, + cooldown, + } => { + edit( + &db, + user_id, + &name, + schedule.as_deref(), + prompt.as_deref(), + description.as_deref(), + timezone.as_deref(), + cooldown, + ) + .await + } + RoutinesCommand::Enable { name } => set_enabled(&db, user_id, &name, true).await, + RoutinesCommand::Disable { name } => set_enabled(&db, user_id, &name, false).await, + RoutinesCommand::Delete { name, yes } => delete(&db, user_id, &name, yes).await, + RoutinesCommand::History { name, limit, json } => { + history(&db, user_id, &name, limit, json).await + } + } +} + +// ── List ──────────────────────────────────────────────────── + +async fn list( + db: &Arc, + user_id: &str, + trigger_filter: Option<&str>, + show_disabled: bool, + json: bool, +) -> anyhow::Result<()> { + let routines = db.list_routines(user_id).await?; + + let filtered: Vec<&Routine> = routines + .iter() + .filter(|r| { + trigger_filter + .map(|t| r.trigger.type_tag() == t) + .unwrap_or(true) + }) + .filter(|r| show_disabled || r.enabled) + .collect(); + + if json { + let items: Vec = filtered + .iter() + .map(|r| { + serde_json::json!({ + "id": r.id.to_string(), + "name": r.name, + "trigger": r.trigger.type_tag(), + "enabled": r.enabled, + "next_fire_at": r.next_fire_at, + "last_run_at": r.last_run_at, + "run_count": r.run_count, + "consecutive_failures": r.consecutive_failures, + }) + }) + .collect(); + println!("{}", serde_json::to_string_pretty(&items)?); + return Ok(()); + } + + if filtered.is_empty() { + if let Some(t) = trigger_filter { + println!("No {t} routines found."); + } else { + println!("No routines found."); + } + return Ok(()); + } + + // Header + println!( + "{:<36} {:<20} {:<8} {:<8} {:<22} {:<22} {:>5}", + "ID", "NAME", "TRIGGER", "STATUS", "NEXT FIRE", "LAST RUN", "RUNS" + ); + println!("{}", "-".repeat(130)); + + for r in &filtered { + let status = if r.enabled { + if r.consecutive_failures > 0 { + format!("err({})", r.consecutive_failures) + } else { + "active".to_string() + } + } else { + "disabled".to_string() + }; + + let next_fire = r + .next_fire_at + .map(format_relative) + .unwrap_or_else(|| "-".to_string()); + + let last_run = r + .last_run_at + .map(format_relative) + .unwrap_or_else(|| "-".to_string()); + + let name = truncate(&r.name, 20); + + println!( + "{:<36} {:<20} {:<8} {:<8} {:<22} {:<22} {:>5}", + r.id, + name, + r.trigger.type_tag(), + status, + next_fire, + last_run, + r.run_count, + ); + } + + println!("\n{} routine(s)", filtered.len()); + Ok(()) +} + +// ── Create ────────────────────────────────────────────────── + +#[allow(clippy::too_many_arguments)] +async fn create( + db: &Arc, + user_id: &str, + name: &str, + schedule: &str, + prompt: &str, + description: &str, + timezone: Option<&str>, + cooldown_secs: u64, + notify_channel: Option, +) -> anyhow::Result<()> { + validate_timezone_arg(timezone)?; + + // Validate the cron expression by computing next fire. + let next_fire = next_cron_fire(schedule, timezone) + .map_err(|e| anyhow::anyhow!("Invalid cron schedule: {e}"))?; + + // Check for name conflict. + if db.get_routine_by_name(user_id, name).await?.is_some() { + anyhow::bail!("Routine '{}' already exists", name); + } + + let now = Utc::now(); + let routine = Routine { + id: Uuid::new_v4(), + name: name.to_string(), + description: description.to_string(), + user_id: user_id.to_string(), + enabled: true, + trigger: Trigger::Cron { + schedule: schedule.to_string(), + timezone: timezone.map(String::from), + }, + action: RoutineAction::Lightweight { + prompt: prompt.to_string(), + context_paths: Vec::new(), + max_tokens: 4096, + use_tools: false, + max_tool_rounds: 0, + }, + guardrails: RoutineGuardrails { + cooldown: std::time::Duration::from_secs(cooldown_secs), + max_concurrent: 1, + dedup_window: None, + }, + notify: NotifyConfig { + channel: notify_channel, + user: user_id.to_string(), + on_attention: true, + on_failure: true, + on_success: false, + }, + last_run_at: None, + next_fire_at: next_fire, + run_count: 0, + consecutive_failures: 0, + state: serde_json::json!({}), + created_at: now, + updated_at: now, + }; + + db.create_routine(&routine).await?; + + println!("Created routine '{}'", name); + println!(" ID: {}", routine.id); + println!(" Schedule: {}", schedule); + if let Some(tz) = timezone { + println!(" Timezone: {}", tz); + } + if let Some(nf) = next_fire { + println!(" Next fire: {}", format_relative(nf)); + } + Ok(()) +} + +// ── Edit ──────────────────────────────────────────────────── + +#[allow(clippy::too_many_arguments)] +async fn edit( + db: &Arc, + user_id: &str, + name: &str, + schedule: Option<&str>, + prompt: Option<&str>, + description: Option<&str>, + timezone: Option<&str>, + cooldown: Option, +) -> anyhow::Result<()> { + let mut routine = require_routine(db, user_id, name).await?; + validate_timezone_arg(timezone)?; + + let mut changed = false; + + // Update schedule if provided (only valid for cron routines). + if let Some(new_schedule) = schedule { + let tz = timezone.or(match &routine.trigger { + Trigger::Cron { timezone, .. } => timezone.as_deref(), + _ => None, + }); + let next_fire = next_cron_fire(new_schedule, tz) + .map_err(|e| anyhow::anyhow!("Invalid cron schedule: {e}"))?; + routine.trigger = Trigger::Cron { + schedule: new_schedule.to_string(), + timezone: tz.map(String::from), + }; + routine.next_fire_at = next_fire; + changed = true; + } else if let Some(tz) = timezone { + // Update only timezone, recompute next fire with existing schedule. + if let Trigger::Cron { ref schedule, .. } = routine.trigger { + let next_fire = next_cron_fire(schedule, Some(tz)) + .map_err(|e| anyhow::anyhow!("Invalid cron schedule: {e}"))?; + routine.trigger = Trigger::Cron { + schedule: schedule.clone(), + timezone: Some(tz.to_string()), + }; + routine.next_fire_at = next_fire; + changed = true; + } else { + anyhow::bail!("Cannot set timezone on non-cron trigger"); + } + } + + if let Some(new_prompt) = prompt { + match &mut routine.action { + RoutineAction::Lightweight { prompt: p, .. } => { + *p = new_prompt.to_string(); + changed = true; + } + RoutineAction::FullJob { description: d, .. } => { + *d = new_prompt.to_string(); + changed = true; + } + } + } + + if let Some(new_desc) = description { + routine.description = new_desc.to_string(); + changed = true; + } + + if let Some(cd) = cooldown { + routine.guardrails.cooldown = std::time::Duration::from_secs(cd); + changed = true; + } + + if !changed { + println!("No changes specified."); + return Ok(()); + } + + routine.updated_at = Utc::now(); + db.update_routine(&routine).await?; + println!("Updated routine '{}'", name); + Ok(()) +} + +// ── Enable / Disable ──────────────────────────────────────── + +async fn set_enabled( + db: &Arc, + user_id: &str, + name: &str, + enabled: bool, +) -> anyhow::Result<()> { + let mut routine = require_routine(db, user_id, name).await?; + + if routine.enabled == enabled { + println!( + "Routine '{}' is already {}", + name, + if enabled { "enabled" } else { "disabled" } + ); + return Ok(()); + } + + routine.enabled = enabled; + + // Recompute next fire when enabling a cron routine. + if enabled + && let Trigger::Cron { + ref schedule, + ref timezone, + } = routine.trigger + { + routine.next_fire_at = next_cron_fire(schedule, timezone.as_deref()) + .map_err(|e| anyhow::anyhow!("Failed to compute next fire for stored schedule: {e}"))?; + } + + routine.updated_at = Utc::now(); + db.update_routine(&routine).await?; + println!( + "{} routine '{}'", + if enabled { "Enabled" } else { "Disabled" }, + name + ); + Ok(()) +} + +// ── Delete ────────────────────────────────────────────────── + +async fn delete( + db: &Arc, + user_id: &str, + name: &str, + skip_confirm: bool, +) -> anyhow::Result<()> { + let routine = require_routine(db, user_id, name).await?; + + if !skip_confirm { + println!("Routine: {}", routine.name); + println!(" ID: {}", routine.id); + println!(" Trigger: {}", routine.trigger.type_tag()); + if let Trigger::Cron { ref schedule, .. } = routine.trigger { + println!("Schedule: {}", schedule); + } + println!(" Runs: {}", routine.run_count); + print!("\nDelete this routine? [y/N] "); + std::io::Write::flush(&mut std::io::stdout())?; + + let mut input = String::new(); + std::io::stdin().read_line(&mut input)?; + if !matches!(input.trim().to_lowercase().as_str(), "y" | "yes") { + println!("Cancelled."); + return Ok(()); + } + } + + let deleted = db.delete_routine(routine.id).await?; + if deleted { + println!("Deleted routine '{}'", name); + } else { + anyhow::bail!("Failed to delete routine '{}'", name); + } + Ok(()) +} + +// ── History ───────────────────────────────────────────────── + +async fn history( + db: &Arc, + user_id: &str, + name: &str, + limit: i64, + json: bool, +) -> anyhow::Result<()> { + let routine = require_routine(db, user_id, name).await?; + + let limit = limit.clamp(1, 50); + let runs = db.list_routine_runs(routine.id, limit).await?; + + if json { + let items: Vec = runs + .iter() + .map(|run| { + serde_json::json!({ + "id": run.id.to_string(), + "status": run.status.to_string(), + "started_at": run.started_at, + "completed_at": run.completed_at, + "result_summary": run.result_summary, + "tokens_used": run.tokens_used, + }) + }) + .collect(); + println!("{}", serde_json::to_string_pretty(&items)?); + return Ok(()); + } + + if runs.is_empty() { + println!("No runs found for routine '{}'", name); + return Ok(()); + } + + println!("Run history for '{}' (last {}):\n", name, runs.len()); + + println!( + "{:<36} {:<8} {:<20} {:<12} SUMMARY", + "RUN ID", "STATUS", "STARTED", "DURATION" + ); + println!("{}", "-".repeat(100)); + + for run in &runs { + let duration = run + .completed_at + .map(|end| { + let secs = (end - run.started_at).num_seconds(); + if secs < 60 { + format!("{}s", secs) + } else { + format!("{}m{}s", secs / 60, secs % 60) + } + }) + .unwrap_or_else(|| "running".to_string()); + + let summary = run + .result_summary + .as_deref() + .map(|s| truncate(s, 40)) + .unwrap_or_else(|| "-".to_string()); + + println!( + "{:<36} {:<8} {:<20} {:<12} {}", + run.id, + run.status, + run.started_at.format("%Y-%m-%d %H:%M:%S"), + duration, + summary, + ); + } + + println!("\n{} run(s) shown", runs.len()); + Ok(()) +} + +// ── Shared lookup ──────────────────────────────────────────── + +/// Look up a routine by name. +async fn require_routine( + db: &Arc, + user_id: &str, + name: &str, +) -> anyhow::Result { + db.get_routine_by_name(user_id, name) + .await? + .ok_or_else(|| anyhow::anyhow!("Routine '{}' not found", name)) +} + +fn validate_timezone_arg(timezone: Option<&str>) -> anyhow::Result<()> { + if let Some(tz) = timezone + && crate::timezone::parse_timezone(tz).is_none() + { + anyhow::bail!("Invalid timezone: '{tz}' is not a valid IANA timezone"); + } + Ok(()) +} + +// ── Helpers ───────────────────────────────────────────────── + +/// Format a datetime relative to now (e.g. "in 2h", "3m ago"). +fn format_relative(dt: DateTime) -> String { + let now = Utc::now(); + let diff = dt.signed_duration_since(now); + let secs = diff.num_seconds(); + + if secs.abs() < 60 { + if secs >= 0 { + "in <1m".to_string() + } else { + "<1m ago".to_string() + } + } else if secs.abs() < 3600 { + let mins = secs.abs() / 60; + if secs >= 0 { + format!("in {}m", mins) + } else { + format!("{}m ago", mins) + } + } else if secs.abs() < 86400 { + let hours = secs.abs() / 3600; + if secs >= 0 { + format!("in {}h", hours) + } else { + format!("{}h ago", hours) + } + } else { + let days = secs.abs() / 86400; + if secs >= 0 { + format!("in {}d", days) + } else { + format!("{}d ago", days) + } + } +} + +/// Truncate a string to a maximum character length. +fn truncate(s: &str, max_chars: usize) -> String { + if s.chars().count() <= max_chars { + s.to_string() + } else { + let truncated: String = s.chars().take(max_chars.saturating_sub(2)).collect(); + format!("{}..", truncated) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn format_relative_future() { + let future = Utc::now() + chrono::Duration::hours(2); + let result = format_relative(future); + assert!( + result.starts_with("in "), + "expected 'in ...' for future time, got: {result}" + ); + } + + #[test] + fn format_relative_past() { + let past = Utc::now() - chrono::Duration::minutes(30); + let result = format_relative(past); + assert!( + result.ends_with(" ago"), + "expected '... ago' for past time, got: {result}" + ); + } + + #[test] + fn format_relative_days() { + let far_future = Utc::now() + chrono::Duration::days(3); + let result = format_relative(far_future); + assert!(result.contains('d'), "expected days in: {result}"); + } + + #[test] + fn truncate_short_string() { + assert_eq!(truncate("hello", 10), "hello"); + } + + #[test] + fn truncate_long_string() { + let result = truncate("hello world", 7); + assert_eq!(result, "hello.."); + } + + #[test] + fn truncate_multibyte_safe() { + // Ensure no panic on multi-byte characters. + let cjk = "你好世界测试"; + let result = truncate(cjk, 4); + assert!(result.ends_with(".."), "got: {result}"); + // Must be valid UTF-8 (would have panicked otherwise). + assert!(result.is_char_boundary(result.len())); + } +} diff --git a/src/cli/snapshots/ironclaw__cli__tests__help_output.snap b/src/cli/snapshots/ironclaw__cli__tests__help_output.snap deleted file mode 100644 index 3c941d88..00000000 --- a/src/cli/snapshots/ironclaw__cli__tests__help_output.snap +++ /dev/null @@ -1,33 +0,0 @@ ---- -source: src/cli/mod.rs -assertion_line: 302 -expression: help ---- -Secure personal AI assistant that protects your data and expands its capabilities - -Usage: ironclaw [OPTIONS] [COMMAND] - -Commands: - run Run the AI agent - onboard Run interactive setup wizard - config Manage app configs - tool Manage WASM tools - registry Browse/install extensions - mcp Manage MCP servers - memory Manage workspace memory - pairing Manage DM pairing - service Manage OS service - doctor Run diagnostics - status Show system status - completion Generate completions - import Import from other AI systems - help Print this message or the help of the given subcommand(s) - -Options: - --cli-only Run in interactive CLI mode only (disable other channels) - --no-db Skip database connection (for testing) - -m, --message Single message mode - send one message and exit - -c, --config Configuration file path (optional, uses env vars by default) - --no-onboard Skip first-run onboarding check - -h, --help Print help (see more with '--help') - -V, --version Print version diff --git a/src/cli/snapshots/ironclaw__cli__tests__help_output_without_import.snap b/src/cli/snapshots/ironclaw__cli__tests__help_output_without_import.snap index c1afbc58..c7d8db13 100644 --- a/src/cli/snapshots/ironclaw__cli__tests__help_output_without_import.snap +++ b/src/cli/snapshots/ironclaw__cli__tests__help_output_without_import.snap @@ -13,6 +13,7 @@ Commands: tool Manage WASM tools registry Browse/install extensions channels Manage channels + routines Manage routines mcp Manage MCP servers memory Manage workspace memory pairing Manage DM pairing diff --git a/src/cli/snapshots/ironclaw__cli__tests__long_help_output.snap b/src/cli/snapshots/ironclaw__cli__tests__long_help_output.snap deleted file mode 100644 index 28e9cb08..00000000 --- a/src/cli/snapshots/ironclaw__cli__tests__long_help_output.snap +++ /dev/null @@ -1,49 +0,0 @@ ---- -source: src/cli/mod.rs -assertion_line: 318 -expression: help ---- -IronClaw is a secure AI assistant. Use 'ironclaw --help' for details. -Examples: - ironclaw run # Start the agent - ironclaw config list # List configs - -Usage: ironclaw [OPTIONS] [COMMAND] - -Commands: - run Run the AI agent - onboard Run interactive setup wizard - config Manage app configs - tool Manage WASM tools - registry Browse/install extensions - mcp Manage MCP servers - memory Manage workspace memory - pairing Manage DM pairing - service Manage OS service - doctor Run diagnostics - status Show system status - completion Generate completions - import Import from other AI systems - help Print this message or the help of the given subcommand(s) - -Options: - --cli-only - Run in interactive CLI mode only (disable other channels) - - --no-db - Skip database connection (for testing) - - -m, --message - Single message mode - send one message and exit - - -c, --config - Configuration file path (optional, uses env vars by default) - - --no-onboard - Skip first-run onboarding check - - -h, --help - Print help (see a summary with '-h') - - -V, --version - Print version diff --git a/src/cli/snapshots/ironclaw__cli__tests__long_help_output_without_import.snap b/src/cli/snapshots/ironclaw__cli__tests__long_help_output_without_import.snap index 7c821c52..fb4ad231 100644 --- a/src/cli/snapshots/ironclaw__cli__tests__long_help_output_without_import.snap +++ b/src/cli/snapshots/ironclaw__cli__tests__long_help_output_without_import.snap @@ -16,6 +16,7 @@ Commands: tool Manage WASM tools registry Browse/install extensions channels Manage channels + routines Manage routines mcp Manage MCP servers memory Manage workspace memory pairing Manage DM pairing diff --git a/src/config/embeddings.rs b/src/config/embeddings.rs index c5a84c00..a1c3ecd7 100644 --- a/src/config/embeddings.rs +++ b/src/config/embeddings.rs @@ -23,6 +23,9 @@ pub struct EmbeddingsConfig { pub ollama_base_url: String, /// Embedding vector dimension. Inferred from the model name when not set explicitly. pub dimension: usize, + /// Custom base URL for OpenAI-compatible embedding providers. + /// When set, overrides the default `https://api.openai.com`. + pub openai_base_url: Option, } impl Default for EmbeddingsConfig { @@ -36,6 +39,7 @@ impl Default for EmbeddingsConfig { model, ollama_base_url: "http://localhost:11434".to_string(), dimension, + openai_base_url: None, } } } @@ -74,6 +78,8 @@ impl EmbeddingsConfig { let enabled = parse_bool_env("EMBEDDING_ENABLED", settings.embeddings.enabled)?; + let openai_base_url = optional_env("EMBEDDING_BASE_URL")?; + Ok(Self { enabled, provider, @@ -81,6 +87,7 @@ impl EmbeddingsConfig { model, ollama_base_url, dimension, + openai_base_url, }) } @@ -130,16 +137,27 @@ impl EmbeddingsConfig { } _ => { if let Some(api_key) = self.openai_api_key() { - tracing::debug!( - "Embeddings enabled via OpenAI (model: {}, dim: {})", - self.model, - self.dimension, - ); - Some(Arc::new(crate::workspace::OpenAiEmbeddings::with_model( + let mut provider = crate::workspace::OpenAiEmbeddings::with_model( api_key, &self.model, self.dimension, - ))) + ); + if let Some(ref base_url) = self.openai_base_url { + tracing::debug!( + "Embeddings enabled via OpenAI (model: {}, base_url: {}, dim: {})", + self.model, + base_url, + self.dimension, + ); + provider = provider.with_base_url(base_url); + } else { + tracing::debug!( + "Embeddings enabled via OpenAI (model: {}, dim: {})", + self.model, + self.dimension, + ); + } + Some(Arc::new(provider)) } else { tracing::warn!("Embeddings configured but OPENAI_API_KEY not set"); None @@ -164,6 +182,7 @@ mod tests { std::env::remove_var("EMBEDDING_PROVIDER"); std::env::remove_var("EMBEDDING_MODEL"); std::env::remove_var("OPENAI_API_KEY"); + std::env::remove_var("EMBEDDING_BASE_URL"); } } @@ -247,4 +266,41 @@ mod tests { std::env::remove_var("EMBEDDING_ENABLED"); } } + + #[test] + fn embedding_base_url_parsed_from_env() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_embedding_env(); + + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::set_var("EMBEDDING_BASE_URL", "https://custom.example.com"); + } + + let settings = Settings::default(); + let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); + assert_eq!( + config.openai_base_url.as_deref(), + Some("https://custom.example.com"), + "EMBEDDING_BASE_URL env var should be parsed into openai_base_url" + ); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::remove_var("EMBEDDING_BASE_URL"); + } + } + + #[test] + fn embedding_base_url_defaults_to_none() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_embedding_env(); + + let settings = Settings::default(); + let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); + assert!( + config.openai_base_url.is_none(), + "openai_base_url should be None when EMBEDDING_BASE_URL is not set" + ); + } } diff --git a/src/config/mod.rs b/src/config/mod.rs index afc54372..34c34423 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -18,6 +18,7 @@ pub mod relay; mod routines; mod safety; mod sandbox; +mod search; mod secrets; mod skills; mod transcription; @@ -44,6 +45,7 @@ pub use self::routines::RoutineConfig; pub use self::safety::SafetyConfig; use self::safety::resolve_safety_config; pub use self::sandbox::{ClaudeCodeConfig, SandboxModeConfig}; +pub use self::search::WorkspaceSearchConfig; pub use self::secrets::SecretsConfig; pub use self::skills::SkillsConfig; pub use self::transcription::TranscriptionConfig; @@ -91,6 +93,7 @@ pub struct Config { pub claude_code: ClaudeCodeConfig, pub skills: SkillsConfig, pub transcription: TranscriptionConfig, + pub search: WorkspaceSearchConfig, pub observability: crate::observability::ObservabilityConfig, /// Channel-relay integration (Slack via external relay service). /// Present only when both `CHANNEL_RELAY_URL` and `CHANNEL_RELAY_API_KEY` are set. @@ -166,6 +169,7 @@ impl Config { ..SkillsConfig::default() }, transcription: TranscriptionConfig::default(), + search: WorkspaceSearchConfig::default(), observability: crate::observability::ObservabilityConfig::default(), relay: None, } @@ -318,6 +322,7 @@ impl Config { claude_code: ClaudeCodeConfig::resolve()?, skills: SkillsConfig::resolve()?, transcription: TranscriptionConfig::resolve(settings)?, + search: WorkspaceSearchConfig::resolve()?, observability: crate::observability::ObservabilityConfig { backend: std::env::var("OBSERVABILITY_BACKEND").unwrap_or_else(|_| "none".into()), }, diff --git a/src/config/search.rs b/src/config/search.rs new file mode 100644 index 00000000..9555fecc --- /dev/null +++ b/src/config/search.rs @@ -0,0 +1,211 @@ +use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::error::ConfigError; +use crate::workspace::FusionStrategy; + +/// Workspace search configuration resolved from environment variables. +#[derive(Debug, Clone)] +pub struct WorkspaceSearchConfig { + /// Fusion strategy: "rrf" or "weighted". + pub fusion_strategy: FusionStrategy, + /// RRF constant k (default 60). + pub rrf_k: u32, + /// FTS weight for fusion. + /// + /// [`Default`] uses 0.5. When the configuration is resolved, per-strategy + /// defaults are applied: 0.5 (RRF) or 0.3 (weighted). + pub fts_weight: f32, + /// Vector weight for fusion. + /// + /// [`Default`] uses 0.5. When the configuration is resolved, per-strategy + /// defaults are applied: 0.5 (RRF) or 0.7 (weighted). + pub vector_weight: f32, +} + +impl Default for WorkspaceSearchConfig { + fn default() -> Self { + Self { + fusion_strategy: FusionStrategy::default(), + rrf_k: 60, + fts_weight: 0.5, + vector_weight: 0.5, + } + } +} + +impl WorkspaceSearchConfig { + pub(crate) fn resolve() -> Result { + let fusion_strategy = match optional_env("SEARCH_FUSION_STRATEGY")? { + Some(s) => match s.to_lowercase().as_str() { + "rrf" => FusionStrategy::Rrf, + "weighted" => FusionStrategy::WeightedScore, + other => { + return Err(ConfigError::InvalidValue { + key: "SEARCH_FUSION_STRATEGY".to_string(), + message: format!("must be 'rrf' or 'weighted', got '{other}'"), + }); + } + }, + None => FusionStrategy::default(), + }; + + let rrf_k = parse_optional_env("SEARCH_RRF_K", 60u32)?; + + // Per-strategy weight defaults: RRF uses 0.5/0.5, weighted uses 0.3/0.7 (vector-biased). + let (default_fts, default_vec) = match fusion_strategy { + FusionStrategy::Rrf => (0.5f32, 0.5f32), + FusionStrategy::WeightedScore => (0.3f32, 0.7f32), + }; + let fts_weight = parse_optional_env("SEARCH_FTS_WEIGHT", default_fts)?; + let vector_weight = parse_optional_env("SEARCH_VECTOR_WEIGHT", default_vec)?; + + if !fts_weight.is_finite() || fts_weight < 0.0 { + return Err(ConfigError::InvalidValue { + key: "SEARCH_FTS_WEIGHT".to_string(), + message: "must be a finite, non-negative float".to_string(), + }); + } + if !vector_weight.is_finite() || vector_weight < 0.0 { + return Err(ConfigError::InvalidValue { + key: "SEARCH_VECTOR_WEIGHT".to_string(), + message: "must be a finite, non-negative float".to_string(), + }); + } + if matches!(fusion_strategy, FusionStrategy::WeightedScore) + && fts_weight == 0.0 + && vector_weight == 0.0 + { + return Err(ConfigError::InvalidValue { + key: "SEARCH_FTS_WEIGHT/SEARCH_VECTOR_WEIGHT".to_string(), + message: "weighted fusion requires at least one non-zero weight".to_string(), + }); + } + + Ok(Self { + fusion_strategy, + rrf_k, + fts_weight, + vector_weight, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::helpers::ENV_MUTEX; + + fn clear_search_env() { + // SAFETY: Only called under ENV_MUTEX in tests. + unsafe { + std::env::remove_var("SEARCH_FUSION_STRATEGY"); + std::env::remove_var("SEARCH_RRF_K"); + std::env::remove_var("SEARCH_FTS_WEIGHT"); + std::env::remove_var("SEARCH_VECTOR_WEIGHT"); + } + } + + #[test] + fn defaults_when_no_env() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + let config = WorkspaceSearchConfig::resolve().expect("should resolve"); + assert_eq!(config.fusion_strategy, FusionStrategy::Rrf); + assert_eq!(config.rrf_k, 60); + assert!((config.fts_weight - 0.5).abs() < 0.001); + assert!((config.vector_weight - 0.5).abs() < 0.001); + } + + #[test] + fn env_overrides() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("SEARCH_FUSION_STRATEGY", "weighted"); + std::env::set_var("SEARCH_RRF_K", "30"); + std::env::set_var("SEARCH_FTS_WEIGHT", "0.9"); + std::env::set_var("SEARCH_VECTOR_WEIGHT", "0.1"); + } + + let config = WorkspaceSearchConfig::resolve().expect("should resolve"); + assert_eq!(config.fusion_strategy, FusionStrategy::WeightedScore); + assert_eq!(config.rrf_k, 30); + assert!((config.fts_weight - 0.9).abs() < 0.001); + assert!((config.vector_weight - 0.1).abs() < 0.001); + + clear_search_env(); + } + + #[test] + fn invalid_strategy_rejected() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("SEARCH_FUSION_STRATEGY", "bm25"); + } + + let result = WorkspaceSearchConfig::resolve(); + assert!(result.is_err()); + + clear_search_env(); + } + + #[test] + fn weighted_strategy_defaults() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("SEARCH_FUSION_STRATEGY", "weighted"); + } + + let config = WorkspaceSearchConfig::resolve().expect("should resolve"); + assert_eq!(config.fusion_strategy, FusionStrategy::WeightedScore); + // Weighted mode should default to 0.3 FTS / 0.7 vector + assert!((config.fts_weight - 0.3).abs() < 0.001); + assert!((config.vector_weight - 0.7).abs() < 0.001); + + clear_search_env(); + } + + #[test] + fn weighted_both_zero_rejected() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("SEARCH_FUSION_STRATEGY", "weighted"); + std::env::set_var("SEARCH_FTS_WEIGHT", "0.0"); + std::env::set_var("SEARCH_VECTOR_WEIGHT", "0.0"); + } + + let result = WorkspaceSearchConfig::resolve(); + assert!(result.is_err()); + + clear_search_env(); + } + + #[test] + fn rrf_both_zero_allowed() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("SEARCH_FTS_WEIGHT", "0.0"); + std::env::set_var("SEARCH_VECTOR_WEIGHT", "0.0"); + } + + // RRF ignores weights, so both=0 is fine + let config = WorkspaceSearchConfig::resolve().expect("should resolve"); + assert_eq!(config.fusion_strategy, FusionStrategy::Rrf); + + clear_search_env(); + } +} diff --git a/src/db/libsql/mod.rs b/src/db/libsql/mod.rs index a5c48f3f..dcc5a8b5 100644 --- a/src/db/libsql/mod.rs +++ b/src/db/libsql/mod.rs @@ -16,6 +16,7 @@ mod workspace; use std::path::Path; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use async_trait::async_trait; use chrono::{DateTime, NaiveDateTime, Utc}; @@ -32,6 +33,8 @@ use crate::workspace::MemoryDocument; use crate::db::libsql_migrations; +static NAIVE_TIMESTAMP_LOGGED: AtomicBool = AtomicBool::new(false); + /// Explicit column list for routines table (matches positional access in `row_to_routine_libsql`). pub(crate) const ROUTINE_COLUMNS: &str = "\ id, name, description, user_id, enabled, \ @@ -163,24 +166,27 @@ impl LibSqlBackend { /// /// Returns an error if none of the formats match. pub(crate) fn parse_timestamp(s: &str) -> Result, String> { + let log_naive_timestamp_once = || { + if !NAIVE_TIMESTAMP_LOGGED.swap(true, Ordering::Relaxed) { + tracing::debug!( + timestamp = %s, + "parsed naive timestamp without timezone; assuming UTC for backward compatibility" + ); + } + }; + // RFC 3339 (our canonical write format) if let Ok(dt) = DateTime::parse_from_rfc3339(s) { return Ok(dt.with_timezone(&Utc)); } // 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::debug!( - timestamp = %s, - "parsed naive timestamp without timezone; assuming UTC for backward compatibility" - ); + log_naive_timestamp_once(); 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::debug!( - timestamp = %s, - "parsed naive timestamp without timezone; assuming UTC for backward compatibility" - ); + log_naive_timestamp_once(); return Ok(ndt.and_utc()); } Err(format!("unparseable timestamp: {:?}", s)) diff --git a/src/db/libsql/workspace.rs b/src/db/libsql/workspace.rs index 19000404..68bd58ba 100644 --- a/src/db/libsql/workspace.rs +++ b/src/db/libsql/workspace.rs @@ -14,7 +14,7 @@ use crate::db::WorkspaceStore; use crate::error::WorkspaceError; use crate::workspace::{ MemoryChunk, MemoryDocument, RankedResult, SearchConfig, SearchResult, WorkspaceEntry, - reciprocal_rank_fusion, + fuse_results, }; use chrono::Utc; @@ -614,6 +614,6 @@ impl WorkspaceStore for LibSqlBackend { ); } - Ok(reciprocal_rank_fusion(fts_results, vector_results, config)) + Ok(fuse_results(fts_results, vector_results, config)) } } diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 2a6cc6d1..6488caa5 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -786,6 +786,19 @@ impl ExtensionManager { Self::validate_extension_name(name)?; let kind = self.determine_installed_kind(name).await?; + // Clean up any in-progress OAuth flows for this extension. + // TCP mode: abort the listener task so port 9876 is freed immediately. + // Gateway mode: remove stale pending flow entries. + if let Some(pending) = self.pending_auth.write().await.remove(name) + && let Some(handle) = pending.task_handle + { + handle.abort(); + } + self.pending_oauth_flows + .write() + .await + .retain(|_, flow| flow.extension_name != name); + match kind { ExtensionKind::McpServer => { // Unregister tools with this server's prefix @@ -819,6 +832,14 @@ impl ExtensionManager { // Unregister from tool registry self.tool_registry.unregister(name).await; + // Evict compiled module from runtime cache so reinstall uses fresh binary + if let Some(ref rt) = self.wasm_tool_runtime { + rt.remove(name).await; + } + + // Clear stale activation errors so reinstall starts clean + self.activation_errors.write().await.remove(name); + // Revoke credential mappings from the shared registry let cap_path = self .wasm_tools_dir @@ -859,6 +880,9 @@ impl ExtensionManager { self.active_channel_names.write().await.remove(name); self.persist_active_channels().await; + // Clear stale activation errors so reinstall starts clean + self.activation_errors.write().await.remove(name); + // Delete channel files let wasm_path = self.wasm_channels_dir.join(format!("{}.wasm", name)); let cap_path = self @@ -2860,6 +2884,17 @@ impl ExtensionManager { }); } + // Check auth status — block activation if required secrets are missing. + // NeedsAuth (OAuth not yet completed) is allowed because configure() loads + // the tool first, then starts the OAuth flow to obtain the token. + let auth_state = self.check_tool_auth_status(name).await; + if auth_state == ToolAuthState::NeedsSetup { + return Err(ExtensionError::ActivationFailed(format!( + "Tool '{}' requires configuration. Use the setup form to provide credentials.", + name + ))); + } + let runtime = self.wasm_tool_runtime.as_ref().ok_or_else(|| { ExtensionError::ActivationFailed("WASM runtime not available".to_string()) })?; @@ -4495,14 +4530,18 @@ mod tests { // available" because the ExtensionManager had `wasm_tool_runtime: None`. /// Build a minimal ExtensionManager suitable for unit tests. - fn make_test_manager( + fn make_test_manager_with_dirs( wasm_runtime: Option>, tools_dir: std::path::PathBuf, + channels_dir: std::path::PathBuf, ) -> crate::extensions::manager::ExtensionManager { use crate::secrets::{InMemorySecretsStore, SecretsCrypto}; use crate::tools::mcp::process::McpProcessManager; use crate::tools::mcp::session::McpSessionManager; + std::fs::create_dir_all(&tools_dir).ok(); + std::fs::create_dir_all(&channels_dir).ok(); + let key = secrecy::SecretString::from(crate::secrets::keychain::generate_master_key_hex()); let crypto = Arc::new(SecretsCrypto::new(key).expect("crypto")); let secrets: Arc = @@ -4517,15 +4556,22 @@ mod tests { tools, None, // hooks wasm_runtime, - tools_dir.clone(), - tools_dir, // channels dir (unused here) - None, // tunnel_url + tools_dir, + channels_dir, + None, // tunnel_url "test".to_string(), None, // db vec![], ) } + fn make_test_manager( + wasm_runtime: Option>, + tools_dir: std::path::PathBuf, + ) -> crate::extensions::manager::ExtensionManager { + make_test_manager_with_dirs(wasm_runtime, tools_dir.clone(), tools_dir) + } + #[tokio::test] async fn test_activate_wasm_tool_with_runtime_passes_runtime_check() { // When the ExtensionManager has a WASM runtime, activation should get @@ -4878,6 +4924,145 @@ mod tests { ); } + #[tokio::test] + async fn test_remove_wasm_tool_clears_pending_oauth_state_and_activation_error() { + let dir = tempfile::tempdir().expect("temp dir"); + let mgr = make_test_manager(None, dir.path().to_path_buf()); + + std::fs::write(dir.path().join("gmail.wasm"), b"fake-tool").expect("write tool"); + + let listener = tokio::spawn(async { + std::future::pending::<()>().await; + }); + let abort_handle = listener.abort_handle(); + mgr.pending_auth.write().await.insert( + "gmail".to_string(), + super::PendingAuth { + _name: "gmail".to_string(), + _kind: ExtensionKind::WasmTool, + created_at: std::time::Instant::now(), + task_handle: Some(listener), + }, + ); + + mgr.activation_errors + .write() + .await + .insert("gmail".to_string(), "cached failure".to_string()); + + let secrets = Arc::clone(&mgr.secrets); + mgr.pending_oauth_flows().write().await.insert( + "gmail-state".to_string(), + crate::cli::oauth_defaults::PendingOAuthFlow { + extension_name: "gmail".to_string(), + display_name: "Gmail".to_string(), + token_url: "https://example.com/token".to_string(), + client_id: "client123".to_string(), + client_secret: None, + redirect_uri: "https://example.com/oauth/callback".to_string(), + code_verifier: None, + access_token_field: "access_token".to_string(), + secret_name: "google_oauth_token".to_string(), + provider: None, + validation_endpoint: None, + scopes: vec![], + user_id: "test".to_string(), + secrets: Arc::clone(&secrets), + sse_sender: None, + gateway_token: None, + resource: None, + client_id_secret_name: None, + created_at: std::time::Instant::now(), + }, + ); + mgr.pending_oauth_flows().write().await.insert( + "other-state".to_string(), + crate::cli::oauth_defaults::PendingOAuthFlow { + extension_name: "web-search".to_string(), + display_name: "Web Search".to_string(), + token_url: "https://example.com/token".to_string(), + client_id: "client456".to_string(), + client_secret: None, + redirect_uri: "https://example.com/oauth/callback".to_string(), + code_verifier: None, + access_token_field: "access_token".to_string(), + secret_name: "other_token".to_string(), + provider: None, + validation_endpoint: None, + scopes: vec![], + user_id: "test".to_string(), + secrets, + sse_sender: None, + gateway_token: None, + resource: None, + client_id_secret_name: None, + created_at: std::time::Instant::now(), + }, + ); + + let result = mgr.remove("gmail").await; + assert!(result.is_ok(), "remove should succeed: {:?}", result.err()); + + tokio::task::yield_now().await; + + assert!( + mgr.pending_auth.read().await.get("gmail").is_none(), + "pending auth entry should be removed" + ); + assert!( + abort_handle.is_finished(), + "pending auth listener should be aborted" + ); + assert!( + !mgr.activation_errors.read().await.contains_key("gmail"), + "stale activation error should be cleared" + ); + + let flows = mgr.pending_oauth_flows().read().await; + assert!( + !flows.contains_key("gmail-state"), + "gateway OAuth flow for removed extension should be cleared" + ); + assert!( + flows.contains_key("other-state"), + "unrelated pending OAuth flows should be retained" + ); + } + + #[tokio::test] + async fn test_remove_wasm_channel_clears_activation_error_and_deletes_files() { + let dir = tempfile::tempdir().expect("temp dir"); + let tools_dir = dir.path().join("tools"); + let channels_dir = dir.path().join("channels"); + let mgr = make_test_manager_with_dirs(None, tools_dir, channels_dir.clone()); + + let wasm_path = channels_dir.join("telegram.wasm"); + let cap_path = channels_dir.join("telegram.capabilities.json"); + std::fs::write(&wasm_path, b"fake-channel").expect("write channel"); + std::fs::write(&cap_path, b"{}").expect("write capabilities"); + + mgr.activation_errors + .write() + .await + .insert("telegram".to_string(), "channel failed".to_string()); + + let result = mgr.remove("telegram").await; + assert!(result.is_ok(), "remove should succeed: {:?}", result.err()); + + assert!( + !mgr.activation_errors.read().await.contains_key("telegram"), + "channel activation error should be cleared on remove" + ); + assert!( + !wasm_path.exists(), + "channel wasm file should be deleted on remove" + ); + assert!( + !cap_path.exists(), + "channel capabilities file should be deleted on remove" + ); + } + #[test] fn test_sanitize_url_with_query_params() { let url = "https://api.example.com/path?api_key=secret123&token=abc"; @@ -5153,7 +5338,6 @@ mod tests { Some("https://my-gateway.example.com/oauth/callback".to_string()), ); } - // ── Regression tests for PR #677 (unify-extension-lifecycle) ───────── #[tokio::test] @@ -5303,7 +5487,6 @@ mod tests { "configure should have stored the relay stream token" ); } - #[test] fn test_validation_failed_is_distinct_error_variant() { // Regression: ValidationFailed must be a distinct error variant so diff --git a/src/main.rs b/src/main.rs index a7bbb295..12a8caf6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -67,6 +67,10 @@ async fn async_main() -> anyhow::Result<()> { ) .await; } + Some(Command::Routines(routines_cmd)) => { + init_cli_tracing(); + return ironclaw::cli::run_routines_cli(routines_cmd, cli.config.as_deref()).await; + } Some(Command::Mcp(mcp_cmd)) => { init_cli_tracing(); return run_mcp_command(*mcp_cmd.clone()).await; @@ -429,9 +433,8 @@ async fn async_main() -> anyhow::Result<()> { "Lifecycle hooks initialized" ); - // Create session manager (shared between agent and web gateway) - let session_manager = - Arc::new(ironclaw::agent::SessionManager::new().with_hooks(components.hooks.clone())); + // Reuse the shared agent session manager prepared by AppBuilder. + let session_manager = Arc::clone(&components.agent_session_manager); // Lazy scheduler slot — filled after Agent::new creates the Scheduler. // Allows CreateJobTool to dispatch local jobs via the Scheduler even though diff --git a/src/service.rs b/src/service.rs index 4249120e..679e6fe2 100644 --- a/src/service.rs +++ b/src/service.rs @@ -65,7 +65,20 @@ fn install_macos() -> Result<()> { let stdout = logs_dir.join("daemon.stdout.log"); let stderr = logs_dir.join("daemon.stderr.log"); - let plist = format!( + let plist = macos_plist_content( + &exe.display().to_string(), + &stdout.display().to_string(), + &stderr.display().to_string(), + ); + + std::fs::write(&file, plist)?; + println!("Installed launchd service: {}", file.display()); + println!(" Start with: ironclaw service start"); + Ok(()) +} + +fn macos_plist_content(exe: &str, stdout: &str, stderr: &str) -> String { + format!( r#" @@ -81,6 +94,11 @@ fn install_macos() -> Result<()> { KeepAlive + EnvironmentVariables + + CLI_ENABLED + false + StandardOutPath {stdout} StandardErrorPath @@ -89,15 +107,10 @@ fn install_macos() -> Result<()> { "#, label = SERVICE_LABEL, - exe = xml_escape(&exe.display().to_string()), - stdout = xml_escape(&stdout.display().to_string()), - stderr = xml_escape(&stderr.display().to_string()), - ); - - std::fs::write(&file, plist)?; - println!("Installed launchd service: {}", file.display()); - println!(" Start with: ironclaw service start"); - Ok(()) + exe = xml_escape(exe), + stdout = xml_escape(stdout), + stderr = xml_escape(stderr), + ) } fn install_linux() -> Result<()> { @@ -356,4 +369,11 @@ mod tests { let s = path.to_string_lossy(); assert!(s.ends_with(".ironclaw/logs"), "unexpected path: {s}"); } + + #[test] + fn macos_plist_sets_cli_enabled_false() { + let plist = macos_plist_content("/tmp/ironclaw", "/tmp/stdout.log", "/tmp/stderr.log"); + assert!(plist.contains("EnvironmentVariables")); + assert!(plist.contains(" CLI_ENABLED\n false")); + } } diff --git a/src/testing/mod.rs b/src/testing/mod.rs index d2078a80..33702e67 100644 --- a/src/testing/mod.rs +++ b/src/testing/mod.rs @@ -1067,6 +1067,8 @@ mod tests { prompt: "Check status".to_string(), context_paths: vec![], max_tokens: 500, + use_tools: false, + max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: std::time::Duration::from_secs(60), @@ -1198,6 +1200,8 @@ mod tests { prompt: "test".to_string(), context_paths: vec![], max_tokens: 100, + use_tools: false, + max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: std::time::Duration::from_secs(0), diff --git a/src/tools/builtin/file.rs b/src/tools/builtin/file.rs index 72e0151c..724b5bae 100644 --- a/src/tools/builtin/file.rs +++ b/src/tools/builtin/file.rs @@ -397,10 +397,6 @@ impl Tool for ListDirTool { false // Directory listings are safe } - fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { - ApprovalRequirement::UnlessAutoApproved - } - fn domain(&self) -> ToolDomain { ToolDomain::Container } diff --git a/src/tools/builtin/http.rs b/src/tools/builtin/http.rs index b1b1994d..4ed2bb0b 100644 --- a/src/tools/builtin/http.rs +++ b/src/tools/builtin/http.rs @@ -31,6 +31,12 @@ const MAX_RESPONSE_SIZE: usize = 5 * 1024 * 1024; /// in memory for LLM context. Matches the WASM attachment size cap. const MAX_SAVE_TO_SIZE: usize = 50 * 1024 * 1024; +/// Default request timeout when the caller does not provide one. +const DEFAULT_TIMEOUT_SECS: u64 = 30; + +/// Maximum allowed request timeout to bound resource usage from LLM-controlled inputs. +const MAX_TIMEOUT_SECS: u64 = 300; + /// Maximum number of redirects to follow for simple GET requests. const MAX_REDIRECTS: usize = 3; @@ -244,43 +250,120 @@ fn is_html_response(headers: &HashMap) -> bool { fn parse_headers_param( headers: Option<&serde_json::Value>, ) -> Result, ToolError> { + fn parse_header_object( + map: &serde_json::Map, + ) -> Result, ToolError> { + let mut out = Vec::with_capacity(map.len()); + for (k, v) in map { + let value = v.as_str().ok_or_else(|| { + ToolError::InvalidParameters(format!("header '{}' must have a string value", k)) + })?; + out.push((k.clone(), value.to_string())); + } + Ok(out) + } + + fn parse_header_array(items: &[serde_json::Value]) -> Result, ToolError> { + let mut out = Vec::with_capacity(items.len()); + for (idx, item) in items.iter().enumerate() { + let obj = item.as_object().ok_or_else(|| { + ToolError::InvalidParameters(format!( + "headers[{}] must be an object with 'name' and 'value'", + idx + )) + })?; + let name = obj.get("name").and_then(|v| v.as_str()).ok_or_else(|| { + ToolError::InvalidParameters(format!("headers[{}].name must be a string", idx)) + })?; + let value = obj.get("value").and_then(|v| v.as_str()).ok_or_else(|| { + ToolError::InvalidParameters(format!("headers[{}].value must be a string", idx)) + })?; + out.push((name.to_string(), value.to_string())); + } + Ok(out) + } + match headers { None => Ok(Vec::new()), - Some(serde_json::Value::Object(map)) => { - let mut out = Vec::with_capacity(map.len()); - for (k, v) in map { - let value = v.as_str().ok_or_else(|| { - ToolError::InvalidParameters(format!("header '{}' must have a string value", k)) - })?; - out.push((k.clone(), value.to_string())); + Some(serde_json::Value::String(raw)) => { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Ok(Vec::new()); } - Ok(out) - } - Some(serde_json::Value::Array(items)) => { - let mut out = Vec::with_capacity(items.len()); - for (idx, item) in items.iter().enumerate() { - let obj = item.as_object().ok_or_else(|| { - ToolError::InvalidParameters(format!( - "headers[{}] must be an object with 'name' and 'value'", - idx - )) - })?; - let name = obj.get("name").and_then(|v| v.as_str()).ok_or_else(|| { - ToolError::InvalidParameters(format!("headers[{}].name must be a string", idx)) - })?; - let value = obj.get("value").and_then(|v| v.as_str()).ok_or_else(|| { - ToolError::InvalidParameters(format!("headers[{}].value must be a string", idx)) - })?; - out.push((name.to_string(), value.to_string())); + let parsed = serde_json::from_str::(trimmed).map_err(|e| { + ToolError::InvalidParameters(format!( + "headers string must contain valid JSON object/array: {}", + e + )) + })?; + match parsed { + serde_json::Value::Object(map) => parse_header_object(&map), + serde_json::Value::Array(items) => parse_header_array(&items), + _ => Err(ToolError::InvalidParameters( + "headers string must decode to a JSON object or array".to_string(), + )), } - Ok(out) } + Some(serde_json::Value::Object(map)) => parse_header_object(map), + Some(serde_json::Value::Array(items)) => parse_header_array(items), Some(_) => Err(ToolError::InvalidParameters( "'headers' must be an object or an array of {name, value}".to_string(), )), } } +fn parse_timeout_secs_param(timeout: Option<&serde_json::Value>) -> Result, ToolError> { + let parsed = match timeout { + None | Some(serde_json::Value::Null) => Ok(None), + Some(serde_json::Value::Number(n)) => n.as_u64().map(Some).ok_or_else(|| { + ToolError::InvalidParameters("timeout_secs must be a non-negative integer".to_string()) + }), + Some(serde_json::Value::String(raw)) => { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Ok(None); + } + let secs = trimmed.parse::().map_err(|_| { + ToolError::InvalidParameters( + "timeout_secs string must contain a non-negative integer".to_string(), + ) + })?; + Ok(Some(secs)) + } + Some(_) => Err(ToolError::InvalidParameters( + "timeout_secs must be an integer".to_string(), + )), + }?; + + if let Some(secs) = parsed + && secs > MAX_TIMEOUT_SECS + { + return Err(ToolError::InvalidParameters(format!( + "timeout_secs must be <= {}", + MAX_TIMEOUT_SECS + ))); + } + + Ok(parsed) +} + +fn parse_save_to_param(save_to: Option<&serde_json::Value>) -> Result, ToolError> { + match save_to { + None | Some(serde_json::Value::Null) => Ok(None), + Some(serde_json::Value::String(path)) => { + let trimmed = path.trim(); + if trimmed.is_empty() { + Ok(None) + } else { + Ok(Some(trimmed.to_string())) + } + } + Some(_) => Err(ToolError::InvalidParameters( + "save_to must be a string".to_string(), + )), + } +} + /// Extract host from URL in params (for approval checks). fn extract_host_from_params(params: &serde_json::Value) -> Option { params @@ -315,7 +398,7 @@ impl Tool for HttpTool { "method": { "type": "string", "enum": ["GET", "POST", "PUT", "DELETE", "PATCH"], - "description": "HTTP method" + "description": "HTTP method (default: GET)" }, "url": { "type": "string", @@ -346,7 +429,7 @@ impl Tool for HttpTool { "description": "Save response body as raw bytes to this file path instead of returning it. Use for binary downloads (images, PDFs, etc.). The path must be under /tmp/." } }, - "required": ["method", "url"] + "required": ["url"] }) } @@ -357,7 +440,8 @@ impl Tool for HttpTool { ) -> Result { let start = std::time::Instant::now(); - let method = require_str(¶ms, "method")?; + let method = params["method"].as_str().unwrap_or("GET"); + let method_upper = method.to_uppercase(); let url = require_str(¶ms, "url")?; let mut parsed_url = validate_url(url)?; @@ -379,6 +463,9 @@ impl Tool for HttpTool { // Parse headers let mut headers_vec = parse_headers_param(params.get("headers"))?; + let timeout_secs = parse_timeout_secs_param(params.get("timeout_secs"))?; + let save_to = parse_save_to_param(params.get("save_to"))?; + let effective_timeout = Duration::from_secs(timeout_secs.unwrap_or(DEFAULT_TIMEOUT_SECS)); // Build request let mut request = match method.to_uppercase().as_str() { @@ -395,6 +482,8 @@ impl Tool for HttpTool { } }; + request = request.timeout(effective_timeout); + // Add headers for (key, value) in &headers_vec { request = request.header(key.as_str(), value.as_str()); @@ -403,7 +492,9 @@ impl Tool for HttpTool { // Add body if present let body_bytes = if let Some(body) = params.get("body") { if let Some(body_str) = body.as_str() { - if let Ok(json_body) = serde_json::from_str::(body_str) { + if body_str.is_empty() { + None + } else if let Ok(json_body) = serde_json::from_str::(body_str) { let bytes = serde_json::to_vec(&json_body).map_err(|e| { ToolError::InvalidParameters(format!("invalid body JSON: {}", e)) })?; @@ -468,7 +559,7 @@ impl Tool for HttpTool { // Build the interceptor request descriptor for recording/replay let intercept_req = crate::llm::recording::HttpExchangeRequest { - method: method.to_uppercase(), + method: method_upper, url: parsed_url.to_string(), headers: headers_vec.clone(), body: body_bytes @@ -510,7 +601,7 @@ impl Tool for HttpTool { let hop_client = build_pinned_client( &hop_host, &hop_addrs, - Duration::from_secs(30), + effective_timeout, reqwest::redirect::Policy::none(), )?; @@ -524,7 +615,7 @@ impl Tool for HttpTool { .await .map_err(|e| { if e.is_timeout() { - ToolError::Timeout(Duration::from_secs(30)) + ToolError::Timeout(effective_timeout) } else { ToolError::ExternalService(e.to_string()) } @@ -588,7 +679,7 @@ impl Tool for HttpTool { } else { let resp = request.send().await.map_err(|e| { if e.is_timeout() { - ToolError::Timeout(Duration::from_secs(30)) + ToolError::Timeout(effective_timeout) } else { ToolError::ExternalService(e.to_string()) } @@ -616,7 +707,7 @@ impl Tool for HttpTool { .collect(); // Use a larger size limit when saving to disk (file downloads) - let saving_to_disk = params.get("save_to").is_some(); + let saving_to_disk = save_to.is_some(); let max_size = if saving_to_disk { MAX_SAVE_TO_SIZE } else { @@ -661,11 +752,11 @@ impl Tool for HttpTool { let body_bytes = bytes::Bytes::from(body); // If save_to is specified, write raw bytes to file and return metadata. - if let Some(save_to) = params.get("save_to").and_then(|v| v.as_str()) { - let save_to_owned = save_to.to_string(); + if let Some(save_to) = save_to { + let saved_to = save_to.clone(); let bytes_clone = body_bytes.clone(); tokio::task::spawn_blocking(move || { - let canonical = validate_save_to_path(&save_to_owned)?; + let canonical = validate_save_to_path(&save_to)?; std::fs::write(&canonical, &bytes_clone).map_err(|e| { ToolError::ExecutionFailed(format!("failed to write file: {}", e)) })?; @@ -676,7 +767,7 @@ impl Tool for HttpTool { .map_err(|e: ToolError| e)?; let result = serde_json::json!({ "status": status, - "saved_to": save_to, + "saved_to": saved_to, "size_bytes": body_bytes.len(), "headers": headers, }); @@ -738,18 +829,22 @@ impl Tool for HttpTool { } fn requires_approval(&self, params: &serde_json::Value) -> ApprovalRequirement { - // 1. Manual auth headers/query params in LLM params - if crate::safety::params_contain_manual_credentials(params) { + let has_credentials = crate::safety::params_contain_manual_credentials(params) + || (self.credential_registry.as_ref().is_some_and(|registry| { + extract_host_from_params(params) + .is_some_and(|host| registry.has_credentials_for_host(&host)) + })); + + if has_credentials { return ApprovalRequirement::Always; } - // 2. Target host has credential mappings (will be auto-injected) - if let Some(ref registry) = self.credential_registry - && let Some(host) = extract_host_from_params(params) - && registry.has_credentials_for_host(&host) - { - return ApprovalRequirement::Always; + + // GET requests (or missing method, since GET is the default) are low-risk + let method = params["method"].as_str().unwrap_or("GET"); + if method.eq_ignore_ascii_case("GET") { + return ApprovalRequirement::Never; } - // Default: outbound HTTP still needs approval unless auto-approved + ApprovalRequirement::UnlessAutoApproved } @@ -887,6 +982,71 @@ mod tests { ); } + #[test] + fn test_parse_headers_param_accepts_stringified_array() { + let headers = + serde_json::json!("[{\"name\":\"Authorization\",\"value\":\"Bearer token\"}]"); + let parsed = parse_headers_param(Some(&headers)).unwrap(); + assert_eq!( + parsed, + vec![("Authorization".to_string(), "Bearer token".to_string())] + ); + } + + #[test] + fn test_parse_headers_param_rejects_double_string_encoding() { + let headers = serde_json::json!("\"hello\""); + let err = parse_headers_param(Some(&headers)).unwrap_err(); + assert!( + err.to_string() + .contains("headers string must decode to a JSON object or array"), + "unexpected error: {}", + err + ); + } + + #[test] + fn test_parse_timeout_secs_param_accepts_string_integer() { + let timeout = serde_json::json!("30"); + assert_eq!(parse_timeout_secs_param(Some(&timeout)).unwrap(), Some(30)); + } + + #[test] + fn test_parse_timeout_secs_param_treats_empty_string_as_none() { + let timeout = serde_json::json!(""); + assert_eq!(parse_timeout_secs_param(Some(&timeout)).unwrap(), None); + } + + #[test] + fn test_parse_timeout_secs_param_rejects_value_above_cap() { + let timeout = serde_json::json!(MAX_TIMEOUT_SECS + 1); + let err = parse_timeout_secs_param(Some(&timeout)).unwrap_err(); + assert!( + err.to_string() + .contains(&format!("timeout_secs must be <= {}", MAX_TIMEOUT_SECS)), + "unexpected error: {}", + err + ); + } + + #[test] + fn test_parse_timeout_secs_param_rejects_string_value_above_cap() { + let timeout = serde_json::json!((MAX_TIMEOUT_SECS + 1).to_string()); + let err = parse_timeout_secs_param(Some(&timeout)).unwrap_err(); + assert!( + err.to_string() + .contains(&format!("timeout_secs must be <= {}", MAX_TIMEOUT_SECS)), + "unexpected error: {}", + err + ); + } + + #[test] + fn test_parse_save_to_param_treats_empty_string_as_none() { + let save_to = serde_json::json!(""); + assert_eq!(parse_save_to_param(Some(&save_to)).unwrap(), None); + } + #[test] fn test_http_tool_schema_body_is_freeform() { let schema = HttpTool::new().parameters_schema(); @@ -907,12 +1067,22 @@ mod tests { // ── Approval requirement tests ────────────────────────────────────── #[test] - fn test_no_auth_headers_returns_unless_auto_approved() { + fn test_get_no_auth_headers_returns_never() { let tool = HttpTool::new(); let params = serde_json::json!({ "method": "GET", "url": "https://api.example.com/data" }); + assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Never); + } + + #[test] + fn test_post_no_auth_headers_returns_unless_auto_approved() { + let tool = HttpTool::new(); + let params = serde_json::json!({ + "method": "POST", + "url": "https://api.example.com/data" + }); assert_eq!( tool.requires_approval(¶ms), ApprovalRequirement::UnlessAutoApproved @@ -996,21 +1166,18 @@ mod tests { } #[test] - fn test_non_auth_headers_return_unless_auto_approved() { + fn test_get_non_auth_headers_return_never() { let tool = HttpTool::new(); let params = serde_json::json!({ "method": "GET", "url": "https://example.com", "headers": {"Content-Type": "application/json", "Accept": "text/html"} }); - assert_eq!( - tool.requires_approval(¶ms), - ApprovalRequirement::UnlessAutoApproved - ); + assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Never); } #[test] - fn test_empty_headers_return_unless_auto_approved() { + fn test_get_empty_headers_return_never() { let tool = HttpTool::new(); // Empty object @@ -1019,10 +1186,7 @@ mod tests { "url": "https://example.com", "headers": {} }); - assert_eq!( - tool.requires_approval(¶ms), - ApprovalRequirement::UnlessAutoApproved - ); + assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Never); // Empty array let params = serde_json::json!({ @@ -1030,10 +1194,7 @@ mod tests { "url": "https://example.com", "headers": [] }); - assert_eq!( - tool.requires_approval(¶ms), - ApprovalRequirement::UnlessAutoApproved - ); + assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Never); } // ── Credential registry approval tests ───────────────────────────── @@ -1063,7 +1224,7 @@ mod tests { } #[test] - fn test_host_without_credential_mapping_returns_unless_auto_approved() { + fn test_get_host_without_credential_mapping_returns_never() { use crate::tools::wasm::SharedCredentialRegistry; let registry = Arc::new(SharedCredentialRegistry::new()); @@ -1075,10 +1236,7 @@ mod tests { "method": "GET", "url": "https://api.example.com/data" }); - assert_eq!( - tool.requires_approval(¶ms), - ApprovalRequirement::UnlessAutoApproved - ); + assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Never); } #[test] @@ -1119,6 +1277,25 @@ mod tests { assert_eq!(extract_host_from_params(¶ms), None); } + #[test] + fn test_requires_approval_with_stringified_http_params() { + use crate::tools::wasm::SharedCredentialRegistry; + + let tool = HttpTool::new().with_credentials( + Arc::new(SharedCredentialRegistry::new()), + Arc::new(test_secrets_store()), + ); + let req = serde_json::json!({ + "body": "", + "headers": "[]", + "method": "GET", + "save_to": "", + "timeout_secs": "30", + "url": "https://r.jina.ai/http://news.baidu.com/" + }); + let _ = tool.requires_approval(&req); + } + // ── DNS pinning tests ───────────────────────────────────────────── #[tokio::test] diff --git a/src/tools/builtin/image_analyze.rs b/src/tools/builtin/image_analyze.rs index b1f8a62f..d6f8f264 100644 --- a/src/tools/builtin/image_analyze.rs +++ b/src/tools/builtin/image_analyze.rs @@ -8,7 +8,7 @@ use secrecy::{ExposeSecret, SecretString}; use crate::context::JobContext; use crate::tools::builtin::path_utils::validate_path; -use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput}; +use crate::tools::tool::{Tool, ToolError, ToolOutput}; /// Tool for analyzing images using a vision-capable model. pub struct ImageAnalyzeTool { @@ -86,10 +86,6 @@ impl Tool for ImageAnalyzeTool { }) } - fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { - ApprovalRequirement::UnlessAutoApproved - } - fn requires_sanitization(&self) -> bool { true } @@ -185,6 +181,7 @@ impl Tool for ImageAnalyzeTool { mod tests { use super::super::media_type_from_path; use super::*; + use crate::tools::tool::ApprovalRequirement; use tempfile::TempDir; #[test] @@ -199,7 +196,7 @@ mod tests { } #[test] - fn test_requires_approval_returns_unless_auto_approved() { + fn test_requires_approval_returns_never() { let tool = ImageAnalyzeTool::new( "https://api.example.com".to_string(), "test-key".to_string(), @@ -208,7 +205,7 @@ mod tests { ); assert_eq!( tool.requires_approval(&serde_json::json!({})), - ApprovalRequirement::UnlessAutoApproved + ApprovalRequirement::Never ); } diff --git a/src/tools/builtin/image_edit.rs b/src/tools/builtin/image_edit.rs index 818454cc..36c2d90d 100644 --- a/src/tools/builtin/image_edit.rs +++ b/src/tools/builtin/image_edit.rs @@ -7,7 +7,7 @@ use secrecy::{ExposeSecret, SecretString}; use crate::context::JobContext; use crate::tools::builtin::path_utils::validate_path; -use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput}; +use crate::tools::tool::{Tool, ToolError, ToolOutput}; /// Tool for editing images using an AI image editing API. pub struct ImageEditTool { @@ -85,10 +85,6 @@ impl Tool for ImageEditTool { }) } - fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { - ApprovalRequirement::UnlessAutoApproved - } - fn requires_sanitization(&self) -> bool { false } @@ -266,6 +262,7 @@ impl ImageEditTool { #[cfg(test)] mod tests { use super::*; + use crate::tools::tool::ApprovalRequirement; use tempfile::TempDir; #[test] @@ -280,7 +277,7 @@ mod tests { assert!(!tool.requires_sanitization()); assert_eq!( tool.requires_approval(&serde_json::json!({})), - ApprovalRequirement::UnlessAutoApproved + ApprovalRequirement::Never ); } diff --git a/src/tools/builtin/image_gen.rs b/src/tools/builtin/image_gen.rs index c87b10d7..a9cc98e9 100644 --- a/src/tools/builtin/image_gen.rs +++ b/src/tools/builtin/image_gen.rs @@ -5,7 +5,6 @@ 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. @@ -87,10 +86,6 @@ impl Tool for ImageGenerateTool { }) } - fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { - ApprovalRequirement::UnlessAutoApproved - } - fn requires_sanitization(&self) -> bool { false } @@ -186,6 +181,7 @@ impl Tool for ImageGenerateTool { #[cfg(test)] mod tests { use super::*; + use crate::tools::tool::ApprovalRequirement; #[test] fn test_tool_metadata() { @@ -197,7 +193,7 @@ mod tests { assert_eq!(tool.name(), "image_generate"); assert_eq!( tool.requires_approval(&serde_json::json!({})), - ApprovalRequirement::UnlessAutoApproved + ApprovalRequirement::Never ); let schema = tool.parameters_schema(); diff --git a/src/tools/builtin/memory.rs b/src/tools/builtin/memory.rs index 71fe8a3b..f1f84684 100644 --- a/src/tools/builtin/memory.rs +++ b/src/tools/builtin/memory.rs @@ -12,6 +12,7 @@ //! Use `memory_write` to persist important facts that should be remembered //! across sessions. +use std::path::Path; use std::sync::Arc; use async_trait::async_trait; @@ -26,6 +27,28 @@ use crate::workspace::{Workspace, paths}; const PROTECTED_IDENTITY_FILES: &[&str] = &[paths::IDENTITY, paths::SOUL, paths::AGENTS, paths::USER]; +/// Detect paths that are clearly local filesystem references, not workspace-memory docs. +/// +/// Examples: +/// - `/Users/.../file.md` (Unix absolute) +/// - `C:\Users\...` or `D:/work/...` (Windows absolute) +/// - `~/notes.md` (home expansion shorthand) +fn looks_like_filesystem_path(path: &str) -> bool { + if path.is_empty() { + return false; + } + + if Path::new(path).is_absolute() || path.starts_with("~/") { + return true; + } + + let bytes = path.as_bytes(); + bytes.len() >= 3 + && bytes[0].is_ascii_alphabetic() + && bytes[1] == b':' + && (bytes[2] == b'\\' || bytes[2] == b'/') +} + /// Tool for searching workspace memory. /// /// Performs hybrid search (FTS + semantic) across all memory documents. @@ -143,7 +166,8 @@ impl Tool for MemoryWriteTool { be remembered across sessions. Targets: 'memory' for curated long-term facts, \ 'daily_log' for timestamped session notes, 'heartbeat' for the periodic \ checklist (HEARTBEAT.md), 'bootstrap' to clear the first-run ritual file, \ - or provide a custom path for arbitrary file creation." + or provide a custom workspace path for arbitrary file creation. \ + Never pass absolute filesystem paths like '/Users/...' or 'C:\\...'." } fn parameters_schema(&self) -> serde_json::Value { @@ -183,6 +207,14 @@ impl Tool for MemoryWriteTool { .and_then(|v| v.as_str()) .unwrap_or("daily_log"); + if looks_like_filesystem_path(target) { + return Err(ToolError::InvalidParameters(format!( + "'{}' looks like a local filesystem path. memory_write only works with workspace-memory paths. \ + Use write_file for filesystem writes. For opening files in an editor, use shell with: open \"\".", + target + ))); + } + // Bootstrap target: clear BOOTSTRAP.md to mark first-run ritual complete. // Handled early because it accepts empty content (unlike other targets). if target == "bootstrap" { @@ -332,7 +364,8 @@ impl Tool for MemoryReadTool { fn description(&self) -> &str { "Read a file from the workspace memory (database-backed storage). \ Use this to read files shown by memory_tree. NOT for local filesystem files \ - (use read_file for those). Works with identity files, heartbeat checklist, \ + (use read_file for those). Do not pass absolute paths like '/Users/...' or 'C:\\...'. \ + Works with identity files, heartbeat checklist, \ memory, daily logs, or any custom workspace path." } @@ -358,6 +391,14 @@ impl Tool for MemoryReadTool { let path = require_str(¶ms, "path")?; + if looks_like_filesystem_path(path) { + return Err(ToolError::InvalidParameters(format!( + "'{}' looks like a local filesystem path. memory_read only works with workspace-memory paths. \ + Use read_file for filesystem reads. For opening files in an editor, use shell with: open \"\".", + path + ))); + } + let doc = self .workspace .read(path) @@ -498,80 +539,100 @@ impl Tool for MemoryTreeTool { } } -#[cfg(all(test, feature = "postgres"))] +#[cfg(test)] mod tests { use super::*; - fn make_test_workspace() -> Arc { - Arc::new(Workspace::new( - "test_user", - deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new( - tokio_postgres::Config::new(), - tokio_postgres::NoTls, + #[test] + fn detects_filesystem_paths() { + assert!(looks_like_filesystem_path("/Users/nige/file.md")); + assert!(looks_like_filesystem_path("C:\\Users\\nige\\file.md")); + assert!(looks_like_filesystem_path("D:/work/file.md")); + assert!(looks_like_filesystem_path("~/notes.md")); + } + + #[test] + fn allows_workspace_memory_paths() { + assert!(!looks_like_filesystem_path("MEMORY.md")); + assert!(!looks_like_filesystem_path("daily/2026-03-11.md")); + assert!(!looks_like_filesystem_path("projects/alpha/notes.md")); + } + + #[cfg(feature = "postgres")] + mod postgres_schema_tests { + use super::*; + + fn make_test_workspace() -> Arc { + Arc::new(Workspace::new( + "test_user", + deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new( + tokio_postgres::Config::new(), + tokio_postgres::NoTls, + )) + .build() + .unwrap(), )) - .build() - .unwrap(), - )) - } + } - #[test] - fn test_memory_search_schema() { - let workspace = make_test_workspace(); - let tool = MemorySearchTool::new(workspace); + #[test] + fn test_memory_search_schema() { + let workspace = make_test_workspace(); + let tool = MemorySearchTool::new(workspace); - assert_eq!(tool.name(), "memory_search"); - assert!(!tool.requires_sanitization()); + assert_eq!(tool.name(), "memory_search"); + assert!(!tool.requires_sanitization()); - let schema = tool.parameters_schema(); - assert!(schema["properties"]["query"].is_object()); - assert!( - schema["required"] - .as_array() - .unwrap() - .contains(&"query".into()) - ); - } + let schema = tool.parameters_schema(); + assert!(schema["properties"]["query"].is_object()); + assert!( + schema["required"] + .as_array() + .unwrap() + .contains(&"query".into()) + ); + } - #[test] - fn test_memory_write_schema() { - let workspace = make_test_workspace(); - let tool = MemoryWriteTool::new(workspace); + #[test] + fn test_memory_write_schema() { + let workspace = make_test_workspace(); + let tool = MemoryWriteTool::new(workspace); - assert_eq!(tool.name(), "memory_write"); + assert_eq!(tool.name(), "memory_write"); - let schema = tool.parameters_schema(); - assert!(schema["properties"]["content"].is_object()); - assert!(schema["properties"]["target"].is_object()); - assert!(schema["properties"]["append"].is_object()); - } + let schema = tool.parameters_schema(); + assert!(schema["properties"]["content"].is_object()); + assert!(schema["properties"]["target"].is_object()); + assert!(schema["properties"]["append"].is_object()); + } - #[test] - fn test_memory_read_schema() { - let workspace = make_test_workspace(); - let tool = MemoryReadTool::new(workspace); + #[test] + fn test_memory_read_schema() { + let workspace = make_test_workspace(); + let tool = MemoryReadTool::new(workspace); - assert_eq!(tool.name(), "memory_read"); + assert_eq!(tool.name(), "memory_read"); - let schema = tool.parameters_schema(); - assert!(schema["properties"]["path"].is_object()); - assert!( - schema["required"] - .as_array() - .unwrap() - .contains(&"path".into()) - ); - } + let schema = tool.parameters_schema(); + assert!(schema["properties"]["path"].is_object()); + assert!( + schema["required"] + .as_array() + .unwrap() + .contains(&"path".into()) + ); + } - #[test] - fn test_memory_tree_schema() { - let workspace = make_test_workspace(); - let tool = MemoryTreeTool::new(workspace); + #[test] + fn test_memory_tree_schema() { + let workspace = make_test_workspace(); + let tool = MemoryTreeTool::new(workspace); - assert_eq!(tool.name(), "memory_tree"); + assert_eq!(tool.name(), "memory_tree"); - let schema = tool.parameters_schema(); - assert!(schema["properties"]["path"].is_object()); - assert!(schema["properties"]["depth"].is_object()); - assert_eq!(schema["properties"]["depth"]["default"], 1); + let schema = tool.parameters_schema(); + assert!(schema["properties"]["path"].is_object()); + assert!(schema["properties"]["depth"].is_object()); + assert_eq!(schema["properties"]["depth"]["default"], 1); + } } } diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index b52502c9..8ba8e57b 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -15,6 +15,7 @@ pub mod secrets_tools; pub(crate) mod shell; pub mod skill_tools; mod time; +mod tool_info; pub use echo::EchoTool; pub use extension_tools::{ @@ -39,6 +40,7 @@ pub use secrets_tools::{SecretDeleteTool, SecretListTool}; pub use shell::ShellTool; pub use skill_tools::{SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool}; pub use time::TimeTool; +pub use tool_info::ToolInfoTool; mod html_converter; pub mod image_analyze; pub mod image_edit; diff --git a/src/tools/builtin/routine.rs b/src/tools/builtin/routine.rs index 573c3c60..43e2add7 100644 --- a/src/tools/builtin/routine.rs +++ b/src/tools/builtin/routine.rs @@ -104,6 +104,14 @@ impl Tool for RoutineCreateTool { "enum": ["lightweight", "full_job"], "description": "Execution mode: 'lightweight' (single LLM call, default) or 'full_job' (multi-turn with tools)" }, + "use_tools": { + "type": "boolean", + "description": "Enable tool access in lightweight mode (default: false). Only safe tools (no approval required) are available. Ignored for full_job mode." + }, + "max_tool_rounds": { + "type": "integer", + "description": "Max tool call rounds in lightweight mode (default: 3). Only used when use_tools is true." + }, "cooldown_secs": { "type": "integer", "description": "Minimum seconds between fires (default: 300)" @@ -262,11 +270,24 @@ impl Tool for RoutineCreateTool { }) .unwrap_or_default(); + let use_tools = params + .get("use_tools") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + let max_tool_rounds = params + .get("max_tool_rounds") + .and_then(|v| v.as_u64()) + .map(|v| v.clamp(1, crate::agent::routine::MAX_TOOL_ROUNDS_LIMIT as u64) as u32) + .unwrap_or(3); + let action = match action_type { "lightweight" => RoutineAction::Lightweight { prompt: prompt.to_string(), context_paths, max_tokens: 4096, + use_tools, + max_tool_rounds, }, "full_job" => { let tool_permissions = crate::agent::routine::parse_tool_permissions(¶ms); diff --git a/src/tools/builtin/tool_info.rs b/src/tools/builtin/tool_info.rs new file mode 100644 index 00000000..cd94384d --- /dev/null +++ b/src/tools/builtin/tool_info.rs @@ -0,0 +1,183 @@ +//! On-demand tool discovery (like CLI `--help`). +//! +//! Two levels of detail: +//! - Default: name, description, parameter names (compact ~150 bytes) +//! - `include_schema: true`: adds the full typed JSON Schema +//! +//! Keeps the tools array compact (WASM tools use permissive schemas) +//! while allowing precise discovery when needed. + +use std::sync::Weak; + +use async_trait::async_trait; + +use crate::context::JobContext; +use crate::tools::registry::ToolRegistry; +use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str}; + +pub struct ToolInfoTool { + registry: Weak, +} + +impl ToolInfoTool { + pub fn new(registry: Weak) -> Self { + Self { registry } + } +} + +#[async_trait] +impl Tool for ToolInfoTool { + fn name(&self) -> &str { + "tool_info" + } + + fn description(&self) -> &str { + "Get info about any tool: description and parameter names. \ + Set include_schema to true for the full typed parameter schema." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Name of the tool to get info about" + }, + "include_schema": { + "type": "boolean", + "description": "If true, include the full typed JSON Schema for parameters (larger response). Default: false.", + "default": false + } + }, + "required": ["name"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + let name = require_str(¶ms, "name")?; + let include_schema = params + .get("include_schema") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + let registry = self.registry.upgrade().ok_or_else(|| { + ToolError::ExecutionFailed( + "tool registry is no longer available for tool_info".to_string(), + ) + })?; + + let tool = registry.get(name).await.ok_or_else(|| { + ToolError::InvalidParameters(format!("No tool named '{name}' is registered")) + })?; + + let schema = tool.discovery_schema(); + + // Extract just param names from the schema's "properties" keys + let param_names: Vec<&str> = schema + .get("properties") + .and_then(|p| p.as_object()) + .map(|props| props.keys().map(|k| k.as_str()).collect()) + .unwrap_or_default(); + + let mut info = serde_json::json!({ + "name": tool.name(), + "description": tool.description(), + "parameters": param_names, + }); + + if include_schema { + info["schema"] = schema; + } + + Ok(ToolOutput::success(info, start.elapsed())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tools::builtin::EchoTool; + use std::sync::Arc; + + #[tokio::test] + async fn test_tool_info_default_returns_param_names() { + let registry = Arc::new(ToolRegistry::new()); + registry.register(Arc::new(EchoTool)).await; + + let tool = ToolInfoTool::new(Arc::downgrade(®istry)); + let ctx = JobContext::default(); + let result = tool + .execute(serde_json::json!({"name": "echo"}), &ctx) + .await + .unwrap(); + + let info = &result.result; + assert_eq!(info["name"], "echo"); + assert!(!info["description"].as_str().unwrap().is_empty()); + // Default: parameters is an array of names, not the full schema + assert!(info["parameters"].is_array()); + assert!( + info["parameters"] + .as_array() + .unwrap() + .iter() + .any(|v| v.as_str() == Some("message")), + "echo tool should have 'message' parameter: {:?}", + info["parameters"] + ); + // No schema field by default + assert!(info.get("schema").is_none()); + } + + #[tokio::test] + async fn test_tool_info_with_schema() { + let registry = Arc::new(ToolRegistry::new()); + registry.register(Arc::new(EchoTool)).await; + + let tool = ToolInfoTool::new(Arc::downgrade(®istry)); + let ctx = JobContext::default(); + let result = tool + .execute( + serde_json::json!({"name": "echo", "include_schema": true}), + &ctx, + ) + .await + .unwrap(); + + let info = &result.result; + assert_eq!(info["name"], "echo"); + // With include_schema: true, schema field should be present + assert!(info["schema"].is_object()); + assert!(info["schema"]["properties"].is_object()); + } + + #[tokio::test] + async fn test_tool_info_unknown_tool() { + let registry = Arc::new(ToolRegistry::new()); + let tool = ToolInfoTool::new(Arc::downgrade(®istry)); + let ctx = JobContext::default(); + let result = tool + .execute(serde_json::json!({"name": "nonexistent"}), &ctx) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_tool_info_registry_dropped() { + let registry = Arc::new(ToolRegistry::new()); + let tool = ToolInfoTool::new(Arc::downgrade(®istry)); + drop(registry); + + let ctx = JobContext::default(); + let result = tool + .execute(serde_json::json!({"name": "echo"}), &ctx) + .await; + assert!(matches!(result, Err(ToolError::ExecutionFailed(_)))); + } +} diff --git a/src/tools/mcp/auth.rs b/src/tools/mcp/auth.rs index a91cb8fc..81f83832 100644 --- a/src/tools/mcp/auth.rs +++ b/src/tools/mcp/auth.rs @@ -1751,7 +1751,11 @@ mod tests { "https://app.attio.com/oidc/authorize", "test-client", "http://127.0.0.1:9876/callback", - &["mcp".to_string(), "offline_access".to_string(), "openid".to_string()], + &[ + "mcp".to_string(), + "offline_access".to_string(), + "openid".to_string(), + ], Some(&pkce), &extra_params, Some("https://mcp.attio.com/mcp"), diff --git a/src/tools/mod.rs b/src/tools/mod.rs index 833d278b..e49cf396 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -12,6 +12,7 @@ pub mod builtin; pub mod execute; pub mod mcp; pub mod rate_limiter; +pub mod redaction; pub mod schema_validator; pub mod wasm; diff --git a/src/tools/redaction.rs b/src/tools/redaction.rs new file mode 100644 index 00000000..f3bad800 --- /dev/null +++ b/src/tools/redaction.rs @@ -0,0 +1,251 @@ +use serde_json::{Map, Value}; + +const REDACTED: &str = "[REDACTED]"; +const SENSITIVE_EXACT: &[&str] = &[ + "authorization", + "proxy-authorization", + "cookie", + "set-cookie", + "x-api-key", + "api-key", + "api_key", + "access_token", + "refresh_token", + "session_token", + "id_token", + "token", + "password", + "passwd", + "secret", + "client_secret", + "private_key", + "apikey", + "apisecret", +]; + +const SENSITIVE_PARTS: &[&str] = &[ + "password", + "passwd", + "secret", + "credential", + "authorization", + "cookie", + "apikey", + "apisecret", +]; +const TOKEN_PARTS: &[&str] = &["token", "jwt"]; +const KEY_PARTS: &[&str] = &["key"]; +const CONTEXT_PARTS: &[&str] = &[ + "auth", + "oauth", + "authorization", + "api", + "access", + "refresh", + "session", + "bearer", + "private", + "client", + "id", + "app", + "user", + "application", + "account", +]; + +fn split_camel_case_key_parts(key: &str) -> Vec { + if key.is_empty() { + return Vec::new(); + } + + let chars: Vec = key.chars().collect(); + let mut parts = Vec::new(); + let mut start = 0; + + for i in 1..chars.len() { + let prev = chars[i - 1]; + let cur = chars[i]; + let next = chars.get(i + 1).copied(); + + let boundary = (prev.is_ascii_lowercase() && cur.is_ascii_uppercase()) + || (prev.is_ascii_alphabetic() && cur.is_ascii_digit()) + || (prev.is_ascii_digit() && cur.is_ascii_alphabetic()) + || (prev.is_ascii_uppercase() + && cur.is_ascii_uppercase() + && next.map(|n| n.is_ascii_lowercase()).unwrap_or(false)); + + if boundary { + parts.push(chars[start..i].iter().collect::()); + start = i; + } + } + + parts.push(chars[start..].iter().collect::()); + parts +} + +fn tokenize_key_parts(key: &str) -> Vec { + let mut parts = Vec::new(); + + for segment in key.split(|c: char| !c.is_ascii_alphanumeric()) { + if segment.is_empty() { + continue; + } + + parts.extend(split_camel_case_key_parts(segment)); + } + + parts.into_iter().map(|p| p.to_ascii_lowercase()).collect() +} + +fn has_exact(parts: &[String], candidates: &[&str]) -> bool { + parts + .iter() + .any(|part| candidates.iter().any(|candidate| part == candidate)) +} + +fn has_candidate_or_numbered_variant(parts: &[String], candidates: &[&str]) -> bool { + parts.iter().any(|part| { + candidates.iter().any(|candidate| { + if part == candidate { + return true; + } + let Some(suffix) = part.strip_prefix(candidate) else { + return false; + }; + !suffix.is_empty() && suffix.chars().all(|c| c.is_ascii_digit()) + }) + }) +} + +fn has_contextual_suffix(parts: &[String], candidates: &[&str]) -> bool { + parts.iter().any(|part| { + candidates.iter().any(|candidate| { + let Some(prefix) = part.strip_suffix(candidate) else { + return false; + }; + !prefix.is_empty() && CONTEXT_PARTS.contains(&prefix) + }) + }) +} + +fn is_sensitive_key(key: &str) -> bool { + let lower = key.to_ascii_lowercase(); + if SENSITIVE_EXACT.contains(&lower.as_str()) { + return true; + } + + let parts = tokenize_key_parts(key); + if parts.is_empty() { + return false; + } + + if has_candidate_or_numbered_variant(&parts, SENSITIVE_PARTS) { + return true; + } + + let has_token = has_candidate_or_numbered_variant(&parts, TOKEN_PARTS); + let has_key = has_candidate_or_numbered_variant(&parts, KEY_PARTS); + + if has_token && has_key { + return true; + } + + if has_contextual_suffix(&parts, TOKEN_PARTS) || has_contextual_suffix(&parts, KEY_PARTS) { + return true; + } + + let has_context = has_exact(&parts, CONTEXT_PARTS); + has_context && (has_token || has_key) +} + +fn redact_in_place(value: &mut Value) { + match value { + Value::Object(map) => redact_object(map), + Value::Array(items) => { + for item in items { + redact_in_place(item); + } + } + _ => {} + } +} + +fn redact_object(map: &mut Map) { + for (key, val) in map { + if is_sensitive_key(key) { + *val = Value::String(REDACTED.to_string()); + } else { + redact_in_place(val); + } + } +} + +pub fn redact_sensitive_json(value: &Value) -> Value { + let mut cloned = value.clone(); + redact_in_place(&mut cloned); + cloned +} + +#[cfg(test)] +mod tests { + use super::{is_sensitive_key, redact_sensitive_json}; + + #[test] + fn redacts_exact_sensitive_keys() { + let input = serde_json::json!({ + "headers": { + "Authorization": "Bearer abc", + "x-api-key": "k-123", + "content-type": "application/json" + }, + "password": "p@ss" + }); + let out = redact_sensitive_json(&input); + assert_eq!(out["headers"]["Authorization"], "[REDACTED]"); + assert_eq!(out["headers"]["x-api-key"], "[REDACTED]"); + assert_eq!(out["headers"]["content-type"], "application/json"); + assert_eq!(out["password"], "[REDACTED]"); + } + + #[test] + fn redacts_nested_sensitive_keys() { + let input = serde_json::json!({ + "body": { + "clientSecret": "xyz", + "nested": [{"authToken": "123"}, {"query": "ok"}] + } + }); + let out = redact_sensitive_json(&input); + assert_eq!(out["body"]["clientSecret"], "[REDACTED]"); + assert_eq!(out["body"]["nested"][0]["authToken"], "[REDACTED]"); + assert_eq!(out["body"]["nested"][1]["query"], "ok"); + } + + #[test] + fn does_not_over_redact_common_non_sensitive_keys() { + assert!(!is_sensitive_key("author")); + assert!(!is_sensitive_key("authorize_user")); + assert!(!is_sensitive_key("token_count")); + assert!(!is_sensitive_key("tokenize")); + assert!(!is_sensitive_key("oauth_redirect_uri")); + } + + #[test] + fn still_redacts_expected_token_keys() { + assert!(is_sensitive_key("auth_token")); + assert!(is_sensitive_key("oauth_token")); + assert!(is_sensitive_key("accessToken")); + assert!(is_sensitive_key("apiKey")); + assert!(is_sensitive_key("token_key")); + assert!(is_sensitive_key("appTokenKey")); + assert!(is_sensitive_key("userJwt")); + } + + #[test] + fn redacts_lowercase_digit_suffix_segments() { + assert!(is_sensitive_key("password123")); + assert!(is_sensitive_key("secret99")); + assert!(is_sensitive_key("accounttoken2")); + } +} diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 7054eea3..754869c8 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -23,7 +23,7 @@ use crate::tools::builtin::{ ToolUpgradeTool, WriteFileTool, }; use crate::tools::rate_limiter::RateLimiter; -use crate::tools::tool::{Tool, ToolDomain}; +use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain}; use crate::tools::wasm::{ Capabilities, OAuthRefreshConfig, ResourceLimits, SharedCredentialRegistry, WasmError, WasmStorageError, WasmToolRuntime, WasmToolStore, WasmToolWrapper, @@ -75,6 +75,7 @@ const PROTECTED_TOOL_NAMES: &[&str] = &[ "image_generate", "image_edit", "image_analyze", + "tool_info", ]; /// Registry of available tools. @@ -245,6 +246,17 @@ impl ToolRegistry { tracing::debug!("Registered {} built-in tools", self.count()); } + /// Register the `tool_info` discovery tool. + /// + /// Requires `Arc` so the tool can query the registry for other tools' + /// schemas at runtime. Call after `register_builtin_tools()`. + pub fn register_tool_info(self: &Arc) { + use crate::tools::builtin::ToolInfoTool; + let tool = ToolInfoTool::new(Arc::downgrade(self)); + self.register_sync(Arc::new(tool)); + tracing::debug!("Registered tool_info discovery tool"); + } + /// Register only orchestrator-domain tools (safe for the main process). /// /// This registers tools that don't touch the filesystem or run shell commands: @@ -278,6 +290,38 @@ impl ToolRegistry { .collect() } + /// Get tool definitions excluding specific tools by name. + /// + /// Used by lightweight routines to filter out denylisted and approval-gated tools + /// so the LLM only sees tools it is actually allowed to call. + pub async fn tool_definitions_excluding(&self, deny: &[&str]) -> Vec { + let empty_params = serde_json::Value::Object(serde_json::Map::new()); + let mut defs: Vec = self + .tools + .read() + .await + .values() + .filter(|tool| { + // Exclude denylisted tools + if deny.contains(&tool.name()) { + return false; + } + // Exclude tools that require approval + matches!( + tool.requires_approval(&empty_params), + ApprovalRequirement::Never + ) + }) + .map(|tool| ToolDefinition { + name: tool.name().to_string(), + description: tool.description().to_string(), + parameters: tool.parameters_schema(), + }) + .collect(); + defs.sort_unstable_by(|a, b| a.name.cmp(&b.name)); + defs + } + /// Register development tools for building software. /// /// These tools provide shell access, file operations, and code editing diff --git a/src/tools/tool.rs b/src/tools/tool.rs index b5879e5a..4a0fda8d 100644 --- a/src/tools/tool.rs +++ b/src/tools/tool.rs @@ -336,6 +336,17 @@ pub trait Tool: Send + Sync { None } + /// Full parameter schema for discovery and coercion purposes. + /// + /// Unlike `parameters_schema()` (which may be permissive to keep the tools + /// array compact), this returns the complete typed schema. Used by the + /// `tool_info` built-in and by WASM parameter coercion. + /// + /// Default: delegates to `parameters_schema()`. + fn discovery_schema(&self) -> serde_json::Value { + self.parameters_schema() + } + /// Get the tool schema for LLM function calling. fn schema(&self) -> ToolSchema { ToolSchema { diff --git a/src/tools/wasm/error.rs b/src/tools/wasm/error.rs index 8bbb8202..a0900775 100644 --- a/src/tools/wasm/error.rs +++ b/src/tools/wasm/error.rs @@ -1,7 +1,5 @@ //! WASM sandbox error types. -use std::fmt; - use thiserror::Error; /// Errors that can occur during WASM tool execution. @@ -68,13 +66,13 @@ pub enum WasmError { Timeout(std::time::Duration), /// Component returned an error response. - /// When `hint` is non-empty it carries the tool's description and parameter - /// schema so the LLM can retry with correct arguments. + /// When `hint` is non-empty it points the LLM to `tool_info` so it can + /// fetch the tool's full parameter schema on demand. #[error("Tool error: {message}{}", if hint.is_empty() { String::new() } else { format!("\n\nTool usage hint:\n{hint}") })] ToolReturnedError { /// The error message from the WASM tool. message: String, - /// Optional description + schema hint (empty when unavailable). + /// Optional retry hint (empty when unavailable). hint: String, }, @@ -99,73 +97,9 @@ impl From for crate::tools::ToolError { } } -/// Details about a trap that occurred during execution. -#[derive(Debug, Clone)] -pub struct TrapInfo { - /// Human-readable trap message. - pub message: String, - /// Trap code if available. - pub code: Option, -} - -impl fmt::Display for TrapInfo { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match &self.code { - Some(code) => write!(f, "{}: {}", code, self.message), - None => write!(f, "{}", self.message), - } - } -} - -/// Known trap codes from Wasmtime. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum TrapCode { - /// Out of bounds memory access. - MemoryOutOfBounds, - /// Out of bounds table access. - TableOutOfBounds, - /// Indirect call type mismatch. - IndirectCallToNull, - /// Signature mismatch on indirect call. - BadSignature, - /// Integer overflow. - IntegerOverflow, - /// Integer division by zero. - IntegerDivisionByZero, - /// Invalid conversion to integer. - BadConversionToInteger, - /// Unreachable instruction executed. - UnreachableCodeReached, - /// Call stack exhausted. - StackOverflow, - /// Out of fuel. - OutOfFuel, - /// Unknown trap code. - Unknown, -} - -impl fmt::Display for TrapCode { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let s = match self { - TrapCode::MemoryOutOfBounds => "memory out of bounds", - TrapCode::TableOutOfBounds => "table out of bounds", - TrapCode::IndirectCallToNull => "indirect call to null", - TrapCode::BadSignature => "bad signature", - TrapCode::IntegerOverflow => "integer overflow", - TrapCode::IntegerDivisionByZero => "integer division by zero", - TrapCode::BadConversionToInteger => "bad conversion to integer", - TrapCode::UnreachableCodeReached => "unreachable code reached", - TrapCode::StackOverflow => "stack overflow", - TrapCode::OutOfFuel => "out of fuel", - TrapCode::Unknown => "unknown trap", - }; - write!(f, "{}", s) - } -} - #[cfg(test)] mod tests { - use crate::tools::wasm::error::{TrapCode, TrapInfo, WasmError}; + use crate::tools::wasm::error::WasmError; #[test] fn test_error_display() { @@ -180,17 +114,6 @@ mod tests { assert!(err.to_string().contains("10000000")); } - #[test] - fn test_trap_info_display() { - let info = TrapInfo { - message: "access at offset 0x1000".to_string(), - code: Some(TrapCode::MemoryOutOfBounds), - }; - let s = info.to_string(); - assert!(s.contains("memory out of bounds")); - assert!(s.contains("access at offset")); - } - #[test] fn test_conversion_to_tool_error() { let wasm_err = WasmError::Trapped("test trap".to_string()); @@ -218,12 +141,11 @@ mod tests { fn test_tool_returned_error_with_hint() { let err = WasmError::ToolReturnedError { message: "unknown action: foobar".to_string(), - hint: "Description: Gmail tool\nParameters schema: {\"type\":\"object\"}".to_string(), + hint: "Tip: call tool_info(name: \"gmail\", include_schema: true) for the full parameter schema.".to_string(), }; let display = err.to_string(); assert!(display.contains("unknown action: foobar")); assert!(display.contains("Tool usage hint")); - assert!(display.contains("Gmail tool")); - assert!(display.contains("Parameters schema")); + assert!(display.contains("tool_info")); } } diff --git a/src/tools/wasm/limits.rs b/src/tools/wasm/limits.rs index 237247e9..d537a583 100644 --- a/src/tools/wasm/limits.rs +++ b/src/tools/wasm/limits.rs @@ -67,14 +67,8 @@ pub struct WasmResourceLimiter { memory_used: u64, /// Maximum tables allowed. max_tables: u32, - /// Current table count. - #[allow(dead_code)] // Reserved for table limit enforcement - tables_created: u32, /// Maximum instances allowed. max_instances: u32, - /// Current instance count. - #[allow(dead_code)] // Reserved for instance limit enforcement - instances_created: u32, } impl WasmResourceLimiter { @@ -87,9 +81,7 @@ impl WasmResourceLimiter { memory_limit, memory_used: 0, max_tables: 10, - tables_created: 0, max_instances: 10, // Component model needs multiple instances for WASI - instances_created: 0, } } diff --git a/src/tools/wasm/mod.rs b/src/tools/wasm/mod.rs index 647d2837..1998e801 100644 --- a/src/tools/wasm/mod.rs +++ b/src/tools/wasm/mod.rs @@ -96,7 +96,7 @@ pub(crate) mod storage; mod wrapper; // Core types -pub use error::{TrapCode, TrapInfo, WasmError}; +pub use error::WasmError; pub use host::{HostState, LogEntry, LogLevel}; pub use limits::{ DEFAULT_FUEL_LIMIT, DEFAULT_MEMORY_LIMIT, DEFAULT_TIMEOUT, FuelConfig, ResourceLimits, diff --git a/src/tools/wasm/runtime.rs b/src/tools/wasm/runtime.rs index 7645af3e..02c56f61 100644 --- a/src/tools/wasm/runtime.rs +++ b/src/tools/wasm/runtime.rs @@ -123,7 +123,9 @@ pub struct PreparedModule { pub name: String, /// Tool description (cached from component). pub description: String, - /// Parameter schema JSON (cached from component). + /// Full parameter schema JSON extracted from the component. + /// Used for discovery and coercion, not necessarily for the compact + /// schema advertised in the main tools array. pub schema: serde_json::Value, /// Pre-compiled component (cheaply cloneable via internal Arc). component: wasmtime::component::Component, @@ -265,11 +267,29 @@ impl WasmToolRuntime { let component = wasmtime::component::Component::new(&engine, &wasm_bytes) .map_err(|e| WasmError::CompilationFailed(e.to_string()))?; - // We need to instantiate briefly to extract metadata. - // In a full implementation, we'd use WIT bindgen to get typed access. - // For now, we extract what we can from the component. - let description = extract_tool_description(&engine, &component)?; - let schema = extract_tool_schema(&engine, &component)?; + // Briefly instantiate to extract metadata (description + schema) + // from the tool's exports, analogous to MCP's list_tools(). + let effective_limits = limits.clone().unwrap_or(default_limits.clone()); + let (description, schema) = crate::tools::wasm::wrapper::extract_wasm_metadata( + &engine, + &component, + &effective_limits, + ) + .unwrap_or_else(|e| { + tracing::warn!( + name = %name, + error = %e, + "WASM metadata extraction failed, using fallbacks" + ); + ( + "WASM sandboxed tool".to_string(), + serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": true + }), + ) + }); Ok::<_, WasmError>(PreparedModule { name: name.clone(), @@ -321,41 +341,6 @@ impl WasmToolRuntime { } } -/// Extract tool description from a compiled component. -/// -/// Returns a generic fallback. Callers should prefer loading the description -/// from the sidecar `*.capabilities.json` file and overriding via -/// `WasmToolWrapper::with_description()` or the `WasmToolRegistration::description` field. -fn extract_tool_description( - _engine: &Engine, - _component: &wasmtime::component::Component, -) -> Result { - // WIT bindgen extraction is not yet implemented (see TODO #4 in CLAUDE.md). - // Real descriptions come from the capabilities.json sidecar file, which is - // loaded by the WasmToolLoader and passed as an override at registration time. - Ok("WASM sandboxed tool".to_string()) -} - -/// Extract tool parameter schema from a compiled component. -/// -/// Returns a permissive fallback that accepts any JSON object. Callers should -/// prefer loading the schema from the sidecar `*.capabilities.json` file and -/// overriding via `WasmToolWrapper::with_schema()` or the -/// `WasmToolRegistration::schema` field. -fn extract_tool_schema( - _engine: &Engine, - _component: &wasmtime::component::Component, -) -> Result { - // WIT bindgen extraction is not yet implemented (see TODO #4 in CLAUDE.md). - // Real schemas come from the capabilities.json sidecar file, which is - // loaded by the WasmToolLoader and passed as an override at registration time. - Ok(serde_json::json!({ - "type": "object", - "properties": {}, - "additionalProperties": true - })) -} - impl std::fmt::Debug for WasmToolRuntime { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("WasmToolRuntime") diff --git a/src/tools/wasm/wrapper.rs b/src/tools/wasm/wrapper.rs index c0294c51..479acfa1 100644 --- a/src/tools/wasm/wrapper.rs +++ b/src/tools/wasm/wrapper.rs @@ -464,9 +464,10 @@ pub struct WasmToolWrapper { /// Capabilities to grant to this tool. capabilities: Capabilities, /// Cached description (from PreparedModule or override). + /// Stored without any tool_info hints — hints are composed at display time. description: String, - /// Cached schema (from PreparedModule or override). - schema: serde_json::Value, + /// Compact and discovery schemas for this tool. + schemas: WasmToolSchemas, /// Injected credentials for HTTP requests (e.g., OAuth tokens). /// Keys are placeholder names like "GOOGLE_ACCESS_TOKEN". credentials: HashMap, @@ -477,6 +478,84 @@ pub struct WasmToolWrapper { oauth_refresh: Option, } +#[derive(Debug, Clone)] +struct WasmToolSchemas { + /// Compact schema advertised in the main tools array. + /// + /// This stays permissive by default to avoid serializing full exported + /// WASM schemas on every LLM call. Sidecars can override it explicitly. + advertised: serde_json::Value, + /// Full schema available for discovery and coercion. + /// + /// Seeded from the WASM `schema()` export at registration time, unless a + /// sidecar explicitly overrides it. + discovery: serde_json::Value, +} + +impl WasmToolSchemas { + fn permissive_schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": true + }) + } + + fn is_permissive_schema(schema: &serde_json::Value) -> bool { + schema + .get("properties") + .and_then(|p| p.as_object()) + .is_none_or(|p| p.is_empty()) + } + + fn new(discovery: serde_json::Value) -> Self { + Self { + advertised: Self::permissive_schema(), + discovery, + } + } + + fn with_override(&self, schema: serde_json::Value) -> Self { + Self { + advertised: schema.clone(), + discovery: schema, + } + } + + fn is_advertised_permissive(&self) -> bool { + Self::is_permissive_schema(&self.advertised) + } + + fn advertised(&self) -> serde_json::Value { + self.advertised.clone() + } + + fn discovery(&self) -> serde_json::Value { + self.discovery.clone() + } + + /// Return the best schema available for type coercion. + /// + /// Prefers the discovery schema when it has typed properties. Falls back + /// to the `PreparedModule` schema extracted at load time rather than + /// re-calling the WASM `schema()` export mid-execution, which could + /// interact with mutable linear memory state. + fn effective_for_coercion(&self, prepared_schema: &serde_json::Value) -> serde_json::Value { + if !Self::is_permissive_schema(&self.discovery) { + return self.discovery.clone(); + } + + // Fall back to the load-time extracted schema from PreparedModule. + // This avoids calling schema() on the already-running WASM instance + // where mutable state could produce inconsistent results. + if !Self::is_permissive_schema(prepared_schema) { + return prepared_schema.clone(); + } + + self.discovery.clone() + } +} + impl WasmToolWrapper { /// Create a new WASM tool wrapper. pub fn new( @@ -486,7 +565,7 @@ impl WasmToolWrapper { ) -> Self { Self { description: prepared.description.clone(), - schema: prepared.schema.clone(), + schemas: WasmToolSchemas::new(prepared.schema.clone()), runtime, prepared, capabilities, @@ -504,7 +583,7 @@ impl WasmToolWrapper { /// Override the parameter schema. pub fn with_schema(mut self, schema: serde_json::Value) -> Self { - self.schema = schema; + self.schemas = self.schemas.with_override(schema); self } @@ -615,9 +694,18 @@ impl WasmToolWrapper { } })?; + // Get typed interface — used for execute. + let tool_iface = instance.near_agent_tool(); + + // Determine effective schema for type coercion. + // Prefer the discovery schema when typed; fall back to the load-time + // extracted schema from PreparedModule rather than re-calling the WASM + // export on the already-running instance. + let effective_schema = self.schemas.effective_for_coercion(&self.prepared.schema); + // Coerce string-encoded values to their schema-declared types. // LLMs frequently pass numeric values as strings (e.g. "5" instead of 5). - let params = coerce_params_to_schema(params, &self.schema); + let params = coerce_params_to_schema(params, &effective_schema); // Prepare the request let params_json = serde_json::to_string(¶ms) @@ -629,7 +717,6 @@ impl WasmToolWrapper { }; // Call execute using the generated typed interface - let tool_iface = instance.near_agent_tool(); let response = tool_iface.call_execute(&mut store, &request).map_err(|e| { let error_str = e.to_string(); if error_str.contains("out of fuel") { @@ -644,12 +731,13 @@ impl WasmToolWrapper { // Get logs from host state let logs = store.data_mut().host_state.take_logs(); - // Check for tool-level error — on failure, call the WASM module's - // description() and schema() exports so the LLM can retry with the - // correct parameters without us having to include the (large) schema - // in every request's tools array. + // Check for tool-level error — point the LLM to tool_info for the + // full schema instead of dumping ~3.5KB inline. if let Some(err) = response.error { - let hint = build_tool_hint(tool_iface, &mut store); + let hint = format!( + "Tip: call tool_info(name: \"{}\", include_schema: true) for the full parameter schema.", + self.prepared.name + ); return Err(WasmError::ToolReturnedError { message: err, hint }); } @@ -658,47 +746,55 @@ impl WasmToolWrapper { } } -/// Maximum characters for the description portion of a tool hint. -const HINT_DESC_MAX: usize = 500; -/// Maximum characters for the schema portion of a tool hint. -const HINT_SCHEMA_MAX: usize = 3000; +/// Extract metadata (description + schema) from a WASM tool by briefly +/// instantiating it and calling its `description()` and `schema()` exports. +/// Analogous to MCP's `list_tools()` — discovers tool capabilities at load time. +/// +/// Falls back to generic description and permissive schema on failure. +pub(super) fn extract_wasm_metadata( + engine: &wasmtime::Engine, + component: &wasmtime::component::Component, + limits: &ResourceLimits, +) -> Result<(String, serde_json::Value), WasmError> { + let store_data = StoreData::new( + limits.memory_bytes, + Capabilities::default(), + HashMap::new(), + vec![], + ); + let mut store = Store::new(engine, store_data); -/// Call the WASM module's `description()` and `schema()` exports to build a -/// hint string. Returns an empty string if both calls fail or return empty. -/// Description is capped at [`HINT_DESC_MAX`] chars, schema at -/// [`HINT_SCHEMA_MAX`] chars. -fn build_tool_hint(tool_iface: &wit_tool::Guest, store: &mut Store) -> String { - let desc = tool_iface - .call_description(&mut *store) + // Configure fuel + epoch deadline so extraction can't hang + if let Err(e) = store.set_fuel(limits.fuel) { + tracing::debug!("Fuel not enabled for metadata extraction: {e}"); + } + store.epoch_deadline_trap(); + let ticks = (limits.timeout.as_millis() / EPOCH_TICK_INTERVAL.as_millis()).max(1) as u64; + store.set_epoch_deadline(ticks); + store.limiter(|data| &mut data.limiter); + + // Instantiate with minimal linker + let mut linker = Linker::new(engine); + WasmToolWrapper::add_host_functions(&mut linker)?; + let instance = SandboxedTool::instantiate(&mut store, component, &linker) + .map_err(|e| WasmError::InstantiationFailed(e.to_string()))?; + let tool_iface = instance.near_agent_tool(); + + // Extract description (fall back to generic) + let description = tool_iface + .call_description(&mut store) + .unwrap_or_else(|_| "WASM sandboxed tool".to_string()); + + // Extract and parse schema (fall back to permissive) + let schema = tool_iface + .call_schema(&mut store) .ok() - .unwrap_or_default(); - let schema = tool_iface.call_schema(&mut *store).ok().unwrap_or_default(); - if desc.is_empty() && schema.is_empty() { - return String::new(); - } - let mut hint = String::new(); - if !desc.is_empty() { - hint.push_str("Description: "); - if desc.len() > HINT_DESC_MAX { - let end = crate::util::floor_char_boundary(&desc, HINT_DESC_MAX); - hint.push_str(&desc[..end]); - hint.push('…'); - } else { - hint.push_str(&desc); - } - hint.push('\n'); - } - if !schema.is_empty() { - hint.push_str("Parameters schema: "); - if schema.len() > HINT_SCHEMA_MAX { - let end = crate::util::floor_char_boundary(&schema, HINT_SCHEMA_MAX); - hint.push_str(&schema[..end]); - hint.push('…'); - } else { - hint.push_str(&schema); - } - } - hint + .and_then(|s| serde_json::from_str::(&s).ok()) + .unwrap_or_else(|| { + serde_json::json!({"type": "object", "properties": {}, "additionalProperties": true}) + }); + + Ok((description, schema)) } #[async_trait] @@ -712,7 +808,33 @@ impl Tool for WasmToolWrapper { } fn parameters_schema(&self) -> serde_json::Value { - self.schema.clone() + self.schemas.advertised() + } + + fn discovery_schema(&self) -> serde_json::Value { + self.schemas.discovery() + } + + /// Compose the tool schema for LLM function calling. + /// + /// When the advertised schema is permissive (no typed properties), appends + /// a hint to the description directing the LLM to call `tool_info` for the + /// full parameter schema. This keeps the raw description clean while still + /// guiding the LLM. + fn schema(&self) -> crate::tools::tool::ToolSchema { + let description = if self.schemas.is_advertised_permissive() { + format!( + "{} (call tool_info(name: \"{}\", include_schema: true) for parameter schema)", + self.description, self.prepared.name + ) + } else { + self.description.clone() + }; + crate::tools::tool::ToolSchema { + name: self.prepared.name.clone(), + description, + parameters: self.schemas.advertised(), + } } async fn execute( @@ -749,7 +871,7 @@ impl Tool for WasmToolWrapper { let prepared = Arc::clone(&self.prepared); let capabilities = self.capabilities.clone(); let description = self.description.clone(); - let schema = self.schema.clone(); + let schemas = self.schemas.clone(); let credentials = self.credentials.clone(); // Execute in blocking task with timeout @@ -759,7 +881,7 @@ impl Tool for WasmToolWrapper { prepared, capabilities, description, - schema, + schemas, credentials, secrets_store: None, // Not needed in blocking task oauth_refresh: None, // Already used above for pre-refresh @@ -1232,6 +1354,7 @@ mod tests { TEST_GOOGLE_OAUTH_TOKEN, TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET, test_secrets_store, }; + use crate::tools::tool::Tool; use crate::tools::wasm::capabilities::Capabilities; use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime}; @@ -1246,6 +1369,84 @@ mod tests { assert!(runtime.config().fuel_config.enabled); } + #[tokio::test] + async fn test_advertised_schema_stays_permissive_until_sidecar_override() { + let discovery_schema = serde_json::json!({ + "type": "object", + "properties": { + "query": { "type": "string" }, + "limit": { "type": "integer" } + }, + "required": ["query"] + }); + + let runtime = Arc::new(WasmToolRuntime::new(WasmRuntimeConfig::for_testing()).unwrap()); + let prepared = runtime + .prepare("search", b"\0asm\x0d\0\x01\0", None) + .await + .unwrap(); + let mut wrapper = + super::WasmToolWrapper::new(Arc::clone(&runtime), prepared, Capabilities::default()); + wrapper.schemas = super::WasmToolSchemas::new(discovery_schema.clone()); + wrapper.description = "Search documents".to_string(); + + // Advertised schema stays permissive; discovery holds the typed schema + assert_eq!( + wrapper.parameters_schema(), + serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": true + }) + ); + assert_eq!(wrapper.discovery_schema(), discovery_schema); + + // Raw description is clean — no tool_info hint baked in + assert!(!wrapper.description().contains("tool_info")); + + // But schema() composes the hint at display time when advertised is permissive + let schema = wrapper.schema(); + assert!( + schema.description.contains("tool_info"), + "schema().description should contain tool_info hint: {}", + schema.description + ); + assert!( + schema.description.contains("include_schema: true"), + "hint should mention include_schema: true: {}", + schema.description + ); + + // After sidecar override, both schemas match and hint disappears + let wrapper = wrapper.with_schema(serde_json::json!({ + "type": "object", + "properties": { + "query": { "type": "string" } + }, + "required": ["query"] + })); + + assert_eq!( + wrapper.parameters_schema(), + serde_json::json!({ + "type": "object", + "properties": { + "query": { "type": "string" } + }, + "required": ["query"] + }) + ); + assert_eq!(wrapper.discovery_schema(), wrapper.parameters_schema()); + + // With typed schema, schema() should NOT include tool_info hint + let schema = wrapper.schema(); + assert!( + !schema.description.contains("tool_info"), + "schema().description should not contain tool_info hint when typed: {}", + schema.description + ); + } + #[test] fn test_capabilities_default() { let caps = Capabilities::default(); @@ -1788,6 +1989,23 @@ mod tests { assert_eq!(result["count"], serde_json::json!("not-a-number")); } + /// Regression: permissive fallback schema (empty properties) must NOT coerce. + /// This documents the bug where WASM tools with no sidecar `parameters` field + /// got the permissive fallback, causing coercion to be a no-op and LLM-provided + /// string integers to reach the WASM tool un-coerced. + #[test] + fn test_coerce_noop_with_permissive_schema() { + let permissive = serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": true + }); + let params = serde_json::json!({"query": "test", "count": "10"}); + let result = super::coerce_params_to_schema(params, &permissive); + // With empty properties, no coercion happens — string stays string + assert_eq!(result["count"], serde_json::json!("10")); + } + /// Regression test: leak scan must run on raw headers (before credential /// injection), not after. If it ran post-injection, the host-injected /// Slack bot token (`xoxb-...`) would trigger a Block and reject the diff --git a/src/workspace/embeddings.rs b/src/workspace/embeddings.rs index 42340fcb..e40337eb 100644 --- a/src/workspace/embeddings.rs +++ b/src/workspace/embeddings.rs @@ -60,12 +60,18 @@ pub trait EmbeddingProvider: Send + Sync { } } +/// Default base URL for the OpenAI API. +const OPENAI_API_BASE_URL: &str = "https://api.openai.com"; + /// OpenAI embedding provider using text-embedding-ada-002 or text-embedding-3-small. +/// +/// Supports any OpenAI-compatible embedding endpoint via [`with_base_url`](Self::with_base_url). pub struct OpenAiEmbeddings { client: reqwest::Client, api_key: String, model: String, dimension: usize, + base_url: String, } impl OpenAiEmbeddings { @@ -78,6 +84,7 @@ impl OpenAiEmbeddings { api_key: api_key.into(), model: "text-embedding-3-small".to_string(), dimension: 1536, + base_url: OPENAI_API_BASE_URL.to_string(), } } @@ -88,6 +95,7 @@ impl OpenAiEmbeddings { api_key: api_key.into(), model: "text-embedding-ada-002".to_string(), dimension: 1536, + base_url: OPENAI_API_BASE_URL.to_string(), } } @@ -98,6 +106,7 @@ impl OpenAiEmbeddings { api_key: api_key.into(), model: "text-embedding-3-large".to_string(), dimension: 3072, + base_url: OPENAI_API_BASE_URL.to_string(), } } @@ -112,8 +121,35 @@ impl OpenAiEmbeddings { api_key: api_key.into(), model: model.into(), dimension, + base_url: OPENAI_API_BASE_URL.to_string(), } } + + /// Set a custom base URL for OpenAI-compatible embedding providers. + /// + /// The URL must use `http://` or `https://` scheme. If no scheme is present, + /// `https://` is prepended automatically. Trailing slashes are stripped. + pub fn with_base_url(mut self, base_url: &str) -> Self { + let url = base_url.trim(); + + // Auto-prepend https:// if no scheme is present. + let mut url = if !url.starts_with("http://") && !url.starts_with("https://") { + tracing::debug!( + "No scheme in embedding base URL '{}', prepending https://", + url + ); + format!("https://{url}") + } else { + url.to_string() + }; + + while url.ends_with('/') { + url.pop(); + } + + self.base_url = url; + self + } } #[derive(Debug, Serialize)] @@ -173,9 +209,11 @@ impl EmbeddingProvider for OpenAiEmbeddings { input: texts, }; + let url = format!("{}/v1/embeddings", self.base_url); + let response = self .client - .post("https://api.openai.com/v1/embeddings") + .post(&url) .header("Authorization", format!("Bearer {}", self.api_key)) .json(&request) .send() @@ -575,9 +613,37 @@ mod tests { let provider = OpenAiEmbeddings::new("test-key"); assert_eq!(provider.dimension(), 1536); assert_eq!(provider.model_name(), "text-embedding-3-small"); + assert_eq!(provider.base_url, OPENAI_API_BASE_URL); let provider = OpenAiEmbeddings::large("test-key"); assert_eq!(provider.dimension(), 3072); assert_eq!(provider.model_name(), "text-embedding-3-large"); + assert_eq!(provider.base_url, OPENAI_API_BASE_URL); + } + + #[test] + fn test_openai_with_base_url_valid() { + let provider = + OpenAiEmbeddings::new("test-key").with_base_url("https://custom.example.com"); + assert_eq!(provider.base_url, "https://custom.example.com"); + } + + #[test] + fn test_openai_with_base_url_strips_trailing_slashes() { + let provider = + OpenAiEmbeddings::new("test-key").with_base_url("https://custom.example.com///"); + assert_eq!(provider.base_url, "https://custom.example.com"); + } + + #[test] + fn test_openai_with_base_url_http_scheme() { + let provider = OpenAiEmbeddings::new("test-key").with_base_url("http://localhost:8080"); + assert_eq!(provider.base_url, "http://localhost:8080"); + } + + #[test] + fn test_openai_with_base_url_schemeless_prepends_https() { + let provider = OpenAiEmbeddings::new("test-key").with_base_url("custom.example.com/v1"); + assert_eq!(provider.base_url, "https://custom.example.com/v1"); } } diff --git a/src/workspace/mod.rs b/src/workspace/mod.rs index fa48072b..ad233caf 100644 --- a/src/workspace/mod.rs +++ b/src/workspace/mod.rs @@ -55,7 +55,9 @@ pub use embeddings::{ }; #[cfg(feature = "postgres")] pub use repository::Repository; -pub use search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion}; +pub use search::{ + FusionStrategy, RankedResult, SearchConfig, SearchResult, fuse_results, reciprocal_rank_fusion, +}; use std::sync::Arc; @@ -332,6 +334,8 @@ pub struct Workspace { storage: WorkspaceStorage, /// Embedding provider for semantic search. embeddings: Option>, + /// Default search configuration applied to all queries. + search_defaults: SearchConfig, } impl Workspace { @@ -343,6 +347,7 @@ impl Workspace { agent_id: None, storage: WorkspaceStorage::Repo(Repository::new(pool)), embeddings: None, + search_defaults: SearchConfig::default(), } } @@ -355,6 +360,7 @@ impl Workspace { agent_id: None, storage: WorkspaceStorage::Db(db), embeddings: None, + search_defaults: SearchConfig::default(), } } @@ -370,6 +376,16 @@ impl Workspace { self } + /// Set the default search configuration from workspace search config. + pub fn with_search_config(mut self, config: &crate::config::WorkspaceSearchConfig) -> Self { + self.search_defaults = SearchConfig::default() + .with_fusion_strategy(config.fusion_strategy) + .with_rrf_k(config.rrf_k) + .with_fts_weight(config.fts_weight) + .with_vector_weight(config.vector_weight); + self + } + /// Get the user ID. pub fn user_id(&self) -> &str { &self.user_id @@ -709,13 +725,13 @@ impl Workspace { /// Hybrid search across all memory documents. /// /// Combines full-text search (BM25) with semantic search (vector similarity) - /// using Reciprocal Rank Fusion (RRF). + /// using the configured fusion strategy. pub async fn search( &self, query: &str, limit: usize, ) -> Result, WorkspaceError> { - self.search_with_config(query, SearchConfig::default().with_limit(limit)) + self.search_with_config(query, self.search_defaults.clone().with_limit(limit)) .await } diff --git a/src/workspace/repository.rs b/src/workspace/repository.rs index de8c3169..82e4f949 100644 --- a/src/workspace/repository.rs +++ b/src/workspace/repository.rs @@ -12,7 +12,7 @@ use uuid::Uuid; use crate::error::WorkspaceError; use crate::workspace::document::{MemoryChunk, MemoryDocument, WorkspaceEntry}; -use crate::workspace::search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion}; +use crate::workspace::search::{RankedResult, SearchConfig, SearchResult, fuse_results}; /// Database repository for workspace operations. pub struct Repository { @@ -415,7 +415,7 @@ impl Repository { Vec::new() }; - Ok(reciprocal_rank_fusion(fts_results, vector_results, config)) + Ok(fuse_results(fts_results, vector_results, config)) } /// Full-text search using PostgreSQL ts_rank_cd. diff --git a/src/workspace/search.rs b/src/workspace/search.rs index dff15298..8b78a125 100644 --- a/src/workspace/search.rs +++ b/src/workspace/search.rs @@ -1,17 +1,30 @@ //! Hybrid search combining full-text and semantic search. //! -//! Uses Reciprocal Rank Fusion (RRF) to combine results from: -//! 1. PostgreSQL full-text search (ts_rank_cd) -//! 2. pgvector cosine similarity search +//! Supports two fusion strategies: +//! 1. **RRF** (Reciprocal Rank Fusion) — the default, rank-based method. +//! `score = sum(1 / (k + rank))` for each retrieval method. +//! 2. **WeightedScore** — converts ranks to scores via `1/rank`, combines with +//! configurable weights (`fts_weight * fts_score + vector_weight * vector_score`), +//! then normalizes to \[0,1\] by dividing by the maximum combined score. //! -//! RRF formula: score = sum(1 / (k + rank)) for each retrieval method -//! This is robust to different score scales and produces better results -//! than simple score averaging. +//! Both strategies combine results from: +//! - PostgreSQL / libSQL full-text search +//! - pgvector / libsql_vector cosine similarity search use std::collections::HashMap; use uuid::Uuid; +/// Strategy used to fuse FTS and vector search results. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum FusionStrategy { + /// Reciprocal Rank Fusion (default). Ignores `fts_weight`/`vector_weight`. + #[default] + Rrf, + /// Weighted score fusion using normalized rank-derived scores. + WeightedScore, +} + /// Configuration for hybrid search. #[derive(Debug, Clone)] pub struct SearchConfig { @@ -27,6 +40,16 @@ pub struct SearchConfig { pub min_score: f32, /// Maximum results to fetch from each method before fusion. pub pre_fusion_limit: usize, + /// Fusion strategy to use when combining results. + pub fusion_strategy: FusionStrategy, + /// Weight for FTS results in `WeightedScore` fusion (default 0.5). + /// Ignored by `Rrf` fusion. For env-based config via + /// `WorkspaceSearchConfig::resolve`, defaults are per-strategy. + pub fts_weight: f32, + /// Weight for vector results in `WeightedScore` fusion (default 0.5). + /// Ignored by `Rrf` fusion. For env-based config via + /// `WorkspaceSearchConfig::resolve`, defaults are per-strategy. + pub vector_weight: f32, } impl Default for SearchConfig { @@ -38,6 +61,9 @@ impl Default for SearchConfig { use_vector: true, min_score: 0.0, pre_fusion_limit: 50, + fusion_strategy: FusionStrategy::default(), + fts_weight: 0.5, + vector_weight: 0.5, } } } @@ -74,6 +100,32 @@ impl SearchConfig { self.min_score = score.clamp(0.0, 1.0); self } + + /// Set the fusion strategy. + pub fn with_fusion_strategy(mut self, strategy: FusionStrategy) -> Self { + self.fusion_strategy = strategy; + self + } + + /// Set the FTS weight for `WeightedScore` fusion. + /// + /// Non-finite (NaN, ±inf) or negative values are ignored. + pub fn with_fts_weight(mut self, weight: f32) -> Self { + if weight.is_finite() && weight >= 0.0 { + self.fts_weight = weight; + } + self + } + + /// Set the vector weight for `WeightedScore` fusion. + /// + /// Non-finite (NaN, ±inf) or negative values are ignored. + pub fn with_vector_weight(mut self, weight: f32) -> Self { + if weight.is_finite() && weight >= 0.0 { + self.vector_weight = weight; + } + self + } } /// A search result with hybrid scoring. @@ -87,7 +139,7 @@ pub struct SearchResult { pub chunk_id: Uuid, /// Chunk content. pub content: String, - /// Combined RRF score (0.0-1.0 normalized). + /// Combined fusion score (0.0-1.0 normalized). Strategy-dependent (RRF or WeightedScore). pub score: f32, /// Rank in FTS results (1-based, None if not in FTS results). pub fts_rank: Option, @@ -123,6 +175,22 @@ pub struct RankedResult { pub rank: u32, // 1-based rank } +/// Fuse FTS and vector search results using the strategy specified in `config`. +/// +/// This is the primary entry point for result fusion. Delegates to +/// [`reciprocal_rank_fusion`] or [`weighted_score_fusion`] based on +/// `config.fusion_strategy`. +pub fn fuse_results( + fts_results: Vec, + vector_results: Vec, + config: &SearchConfig, +) -> Vec { + match config.fusion_strategy { + FusionStrategy::Rrf => reciprocal_rank_fusion(fts_results, vector_results, config), + FusionStrategy::WeightedScore => weighted_score_fusion(fts_results, vector_results, config), + } +} + /// Reciprocal Rank Fusion algorithm. /// /// Combines ranked results from multiple retrieval methods using the formula: @@ -235,6 +303,109 @@ pub fn reciprocal_rank_fusion( results } +/// Weighted score fusion. +/// +/// Converts ranks from each method into scores using `1/rank` +/// (so rank 1 → 1.0, rank N → 1/N), then combines them with +/// configurable weights: `fts_weight * fts_score + vector_weight * vector_score`. +/// +/// The combined scores are then normalized to [0,1] by dividing by the +/// maximum score; post-processing (normalization, min_score filter, sort, +/// truncate) matches RRF. +pub fn weighted_score_fusion( + fts_results: Vec, + vector_results: Vec, + config: &SearchConfig, +) -> Vec { + struct ChunkInfo { + document_id: Uuid, + document_path: String, + content: String, + score: f32, + fts_rank: Option, + vector_rank: Option, + } + + let mut chunk_scores: HashMap = HashMap::new(); + + // Process FTS results: score = fts_weight * (1 / rank) + for result in fts_results { + let score = config.fts_weight * (1.0 / result.rank as f32); + chunk_scores + .entry(result.chunk_id) + .and_modify(|info| { + info.score += score; + info.fts_rank = Some(result.rank); + }) + .or_insert(ChunkInfo { + document_id: result.document_id, + document_path: result.document_path, + content: result.content, + score, + fts_rank: Some(result.rank), + vector_rank: None, + }); + } + + // Process vector results: score = vector_weight * (1 / rank) + for result in vector_results { + let score = config.vector_weight * (1.0 / result.rank as f32); + chunk_scores + .entry(result.chunk_id) + .and_modify(|info| { + info.score += score; + info.vector_rank = Some(result.rank); + }) + .or_insert(ChunkInfo { + document_id: result.document_id, + document_path: result.document_path, + content: result.content, + score, + fts_rank: None, + vector_rank: Some(result.rank), + }); + } + + let mut results: Vec = chunk_scores + .into_iter() + .map(|(chunk_id, info)| SearchResult { + document_id: info.document_id, + document_path: info.document_path, + chunk_id, + content: info.content, + score: info.score, + fts_rank: info.fts_rank, + vector_rank: info.vector_rank, + }) + .collect(); + + // Normalize scores to 0-1 range + if let Some(max_score) = results.iter().map(|r| r.score).reduce(f32::max) + && max_score > 0.0 + { + for result in &mut results { + result.score /= max_score; + } + } + + // Filter by minimum score + if config.min_score > 0.0 { + results.retain(|r| r.score >= config.min_score); + } + + // Sort by score descending + results.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + + // Limit results + results.truncate(config.limit); + + results +} + #[cfg(test)] mod tests { use super::*; @@ -457,6 +628,142 @@ mod tests { let vector_only = SearchConfig::default().vector_only(); assert!(!vector_only.use_fts); assert!(vector_only.use_vector); + + let weighted = SearchConfig::default() + .with_fusion_strategy(FusionStrategy::WeightedScore) + .with_fts_weight(0.8) + .with_vector_weight(0.2); + assert_eq!(weighted.fusion_strategy, FusionStrategy::WeightedScore); + assert!((weighted.fts_weight - 0.8).abs() < 0.001); + assert!((weighted.vector_weight - 0.2).abs() < 0.001); + } + + #[test] + fn test_weighted_fusion_basic() { + // With equal weights, a hybrid match should still rank highest. + let config = SearchConfig::default() + .with_fusion_strategy(FusionStrategy::WeightedScore) + .with_fts_weight(1.0) + .with_vector_weight(1.0) + .with_limit(10); + + let chunk1 = Uuid::new_v4(); // In both + let chunk2 = Uuid::new_v4(); // FTS only + let chunk3 = Uuid::new_v4(); // Vector only + let doc = Uuid::new_v4(); + + let fts = vec![make_result(chunk1, doc, 1), make_result(chunk2, doc, 2)]; + let vec_results = vec![make_result(chunk1, doc, 1), make_result(chunk3, doc, 2)]; + + let results = weighted_score_fusion(fts, vec_results, &config); + + assert_eq!(results.len(), 3); + // Hybrid match (chunk1) should be first — it gets score from both + assert_eq!(results[0].chunk_id, chunk1); + assert!(results[0].is_hybrid()); + assert!(results[0].score > results[1].score); + } + + #[test] + fn test_weighted_fusion_fts_boost() { + // High FTS weight should elevate FTS-only results above vector-only. + let config = SearchConfig::default() + .with_fusion_strategy(FusionStrategy::WeightedScore) + .with_fts_weight(2.0) + .with_vector_weight(0.5) + .with_limit(10); + + let chunk_fts = Uuid::new_v4(); // FTS only, rank 2 + let chunk_vec = Uuid::new_v4(); // Vector only, rank 2 + let doc = Uuid::new_v4(); + + let fts = vec![make_result(chunk_fts, doc, 2)]; + let vec_results = vec![make_result(chunk_vec, doc, 2)]; + + let results = weighted_score_fusion(fts, vec_results, &config); + + assert_eq!(results.len(), 2); + // FTS result should rank higher because of the 2.0 weight vs 0.5 + assert_eq!(results[0].chunk_id, chunk_fts); + assert!(results[0].from_fts()); + assert!(!results[0].from_vector()); + } + + #[test] + fn test_weighted_fusion_single_source() { + // Only FTS results — should still work correctly. + let config = SearchConfig::default() + .with_fusion_strategy(FusionStrategy::WeightedScore) + .with_limit(10); + + let chunk1 = Uuid::new_v4(); + let chunk2 = Uuid::new_v4(); + let doc = Uuid::new_v4(); + + let fts = vec![make_result(chunk1, doc, 1), make_result(chunk2, doc, 3)]; + + let results = weighted_score_fusion(fts, Vec::new(), &config); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].chunk_id, chunk1); + assert!(results[0].score > results[1].score); + // Top result should be normalized to 1.0 + assert!((results[0].score - 1.0).abs() < 0.001); + } + + #[test] + fn test_weight_setters_reject_invalid() { + let config = SearchConfig::default(); + let original_fts = config.fts_weight; + let original_vec = config.vector_weight; + + // NaN is ignored + let c = config.clone().with_fts_weight(f32::NAN); + assert!((c.fts_weight - original_fts).abs() < 0.001); + + // Infinity is ignored + let c = config.clone().with_vector_weight(f32::INFINITY); + assert!((c.vector_weight - original_vec).abs() < 0.001); + + // Negative is ignored + let c = config.clone().with_fts_weight(-1.0); + assert!((c.fts_weight - original_fts).abs() < 0.001); + + // Negative infinity is ignored + let c = config.clone().with_vector_weight(f32::NEG_INFINITY); + assert!((c.vector_weight - original_vec).abs() < 0.001); + + // Valid values > 1.0 are accepted (weights don't need to sum to 1.0) + let c = config.clone().with_fts_weight(2.0); + assert!((c.fts_weight - 2.0).abs() < 0.001); + + // Zero is valid + let c = config.clone().with_vector_weight(0.0); + assert!(c.vector_weight.abs() < 0.001); + } + + #[test] + fn test_fuse_results_dispatches_correctly() { + let chunk1 = Uuid::new_v4(); + let doc = Uuid::new_v4(); + + let fts = vec![make_result(chunk1, doc, 1)]; + + // RRF strategy + let rrf_config = SearchConfig::default().with_limit(10); + let rrf_results = fuse_results(fts.clone(), Vec::new(), &rrf_config); + assert_eq!(rrf_results.len(), 1); + + // Weighted strategy + let weighted_config = SearchConfig::default() + .with_fusion_strategy(FusionStrategy::WeightedScore) + .with_limit(10); + let weighted_results = fuse_results(fts, Vec::new(), &weighted_config); + assert_eq!(weighted_results.len(), 1); + + // Both should normalize single result to 1.0 + assert!((rrf_results[0].score - 1.0).abs() < 0.001); + assert!((weighted_results[0].score - 1.0).abs() < 0.001); } // --- Edge case tests --- diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 41a9fd29..9503136d 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -20,9 +20,31 @@ from helpers import AUTH_TOKEN, wait_for_port_line, wait_for_ready # Project root (two levels up from tests/e2e/) ROOT = Path(__file__).resolve().parent.parent.parent +# Git main repo root (for worktree support — WASM build artifacts live +# in the main repo's tools-src/*/target/ and aren't shared across worktrees) +_MAIN_ROOT = None +try: + import subprocess as _sp + _common = _sp.check_output( + ["git", "worktree", "list", "--porcelain"], + cwd=ROOT, text=True, stderr=_sp.DEVNULL, + ) + for line in _common.splitlines(): + if line.startswith("worktree "): + _MAIN_ROOT = Path(line.split(" ", 1)[1]) + break # first entry is always the main worktree +except Exception: + pass + # Temp directory for the libSQL database file (cleaned up automatically) _DB_TMPDIR = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-") +# Temp directories for WASM extensions. These start empty and are populated by +# the install pipeline during tests; fixtures do not pre-populate dev build +# artifacts into them. +_WASM_TOOLS_TMPDIR = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-wasm-tools-") +_WASM_CHANNELS_TMPDIR = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-wasm-channels-") + def _find_free_port() -> int: """Bind to port 0 and return the OS-assigned port.""" @@ -70,7 +92,53 @@ async def mock_llm_server(): @pytest.fixture(scope="session") -async def ironclaw_server(ironclaw_binary, mock_llm_server): +def wasm_tools_dir(_wasm_build_symlinks): + """Empty temp dir for WASM tools. + + Starts empty so the server has no pre-loaded extensions at boot. + The install API (POST /api/extensions/install) downloads and writes + WASM files here; tests exercise the full install pipeline. + + NOTE on capabilities file naming: Cargo builds with underscored stems + (web_search_tool.wasm) but capabilities use hyphens (web-search-tool. + capabilities.json). The loader expects matching stems. If you pre-load + files, rename caps: web-search-tool → web_search_tool. + """ + return str(Path(_WASM_TOOLS_TMPDIR.name)) + + +@pytest.fixture(scope="session", autouse=True) +def _wasm_build_symlinks(): + """Symlink WASM build artifacts from the main repo into the worktree. + + In a git worktree, tools-src/*/target/ directories don't exist because + Cargo build artifacts aren't shared. The install API's source fallback + checks these paths. Symlinking makes the fallback work without rebuilding. + """ + if _MAIN_ROOT is None or _MAIN_ROOT == ROOT: + yield + return + + created = [] + tools_src = ROOT / "tools-src" + main_tools_src = _MAIN_ROOT / "tools-src" + if tools_src.is_dir() and main_tools_src.is_dir(): + for tool_dir in tools_src.iterdir(): + if not tool_dir.is_dir(): + continue + target = tool_dir / "target" + main_target = main_tools_src / tool_dir.name / "target" + if not target.exists() and main_target.is_dir(): + target.symlink_to(main_target) + created.append(target) + yield + for link in created: + if link.is_symlink(): + link.unlink() + + +@pytest.fixture(scope="session") +async def ironclaw_server(ironclaw_binary, mock_llm_server, wasm_tools_dir): """Start the ironclaw gateway. Yields the base URL.""" gateway_port = _find_free_port() env = { @@ -95,8 +163,16 @@ async def ironclaw_server(ironclaw_binary, mock_llm_server): "ROUTINES_ENABLED": "false", "HEARTBEAT_ENABLED": "false", "EMBEDDING_ENABLED": "false", + # WASM tool/channel support + "WASM_ENABLED": "true", + "WASM_TOOLS_DIR": wasm_tools_dir, + "WASM_CHANNELS_DIR": _WASM_CHANNELS_TMPDIR.name, # Prevent onboarding wizard from triggering "ONBOARD_COMPLETED": "true", + # Force gateway OAuth callback mode (non-loopback URL) and point + # token exchange at mock_llm.py so OAuth tests work without Google. + "IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback", + "IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server, } # Forward LLVM coverage instrumentation env vars when present # (allows cargo-llvm-cov to collect profraw data from E2E runs). diff --git a/tests/e2e/helpers.py b/tests/e2e/helpers.py index b6927dce..629205a1 100644 --- a/tests/e2e/helpers.py +++ b/tests/e2e/helpers.py @@ -133,3 +133,32 @@ async def wait_for_port_line(process, pattern: str, *, timeout: float = 60) -> i if match := re.search(pattern, decoded): return int(match.group(1)) raise TimeoutError(f"Port pattern '{pattern}' not found in stdout after {timeout}s") + + +# -- API helpers ----------------------------------------------------------- + +def auth_headers() -> dict[str, str]: + """Return Authorization header dict for authenticated API calls.""" + return {"Authorization": f"Bearer {AUTH_TOKEN}"} + + +async def api_get(base_url: str, path: str, **kwargs) -> httpx.Response: + """Make an authenticated GET request to the ironclaw API.""" + async with httpx.AsyncClient() as client: + return await client.get( + f"{base_url}{path}", + headers=auth_headers(), + timeout=kwargs.pop("timeout", 10), + **kwargs, + ) + + +async def api_post(base_url: str, path: str, **kwargs) -> httpx.Response: + """Make an authenticated POST request to the ironclaw API.""" + async with httpx.AsyncClient() as client: + return await client.post( + f"{base_url}{path}", + headers=auth_headers(), + timeout=kwargs.pop("timeout", 10), + **kwargs, + ) diff --git a/tests/e2e/mock_llm.py b/tests/e2e/mock_llm.py index deb18bd7..0fa0ce9f 100644 --- a/tests/e2e/mock_llm.py +++ b/tests/e2e/mock_llm.py @@ -1,11 +1,16 @@ -"""Mock OpenAI-compatible LLM server for E2E tests.""" +"""Mock OpenAI-compatible LLM server for E2E tests. + +Serves OpenAI-compatible endpoints for chat completions and model listing. +Supports both streaming and non-streaming responses, plus function calling +via TOOL_CALL_PATTERNS. +""" import argparse +import asyncio import json import re import time import uuid - from aiohttp import web CANNED_RESPONSES = [ @@ -13,85 +18,207 @@ CANNED_RESPONSES = [ (re.compile(r"2\s*\+\s*2|two plus two", re.IGNORECASE), "The answer is 4."), (re.compile(r"skill|install", re.IGNORECASE), "I can help you with skills management."), (re.compile(r"html.?test|injection.?test", re.IGNORECASE), - 'Here is some content: and and end of content.'), + 'Here is some content: and ' + ' and end of content.'), ] DEFAULT_RESPONSE = "I understand your request." +TOOL_CALL_PATTERNS = [ + (re.compile(r"echo (.+)", re.IGNORECASE), "echo", lambda m: {"message": m.group(1)}), + (re.compile(r"what time|current time", re.IGNORECASE), "time", lambda _: {"operation": "now"}), +] -def match_response(messages: list[dict]) -> str: - """Find canned response for the last user message.""" + +def _last_user_content(messages: list[dict]) -> str: for msg in reversed(messages): if msg.get("role") == "user": content = msg.get("content", "") - # Handle content that may be a list (multi-modal) if isinstance(content, list): content = " ".join( - part.get("text", "") for part in content if part.get("type") == "text" + p.get("text", "") for p in content if p.get("type") == "text" ) - for pattern, response in CANNED_RESPONSES: - if pattern.search(content): - return response - return DEFAULT_RESPONSE + return content + return "" + + +def match_response(messages: list[dict]) -> str: + content = _last_user_content(messages) + for pattern, response in CANNED_RESPONSES: + if pattern.search(content): + return response return DEFAULT_RESPONSE +def match_tool_call(messages: list[dict], has_tools: bool) -> dict | None: + if not has_tools: + return None + content = _last_user_content(messages) + for pattern, tool_name, args_fn in TOOL_CALL_PATTERNS: + m = pattern.search(content) + if m: + return {"tool_name": tool_name, "arguments": args_fn(m)} + return None + + +def _extract_tool_name(msg: dict) -> str: + """Extract tool name from a message, checking both 'name' field and XML content.""" + name = msg.get("name") + if name: + return name + # ironclaw wraps tool output as + content = msg.get("content", "") + m = re.search(r' dict | None: + """Find a pending tool result that appears after the last user message. + + Only returns a tool result if it's a fresh result the agent is waiting + for the LLM to summarize (i.e., it follows the most recent user message). + This prevents stale tool results from earlier conversation turns from + being re-processed. + """ + # Find the position of the last user message + last_user_idx = -1 + for i in range(len(messages) - 1, -1, -1): + if messages[i].get("role") == "user": + last_user_idx = i + break + + # Only look for tool results after the last user message + for i in range(len(messages) - 1, last_user_idx, -1): + if messages[i].get("role") == "tool": + return {"name": _extract_tool_name(messages[i]), + "content": messages[i].get("content", "")} + return None + + +def _make_base(completion_id: str) -> dict: + return {"id": completion_id, "object": "chat.completion.chunk", + "created": int(time.time()), "model": "mock-model"} + + +async def _send_sse(resp: web.StreamResponse, data: dict): + await resp.write(f"data: {json.dumps(data)}\n\n".encode()) + + async def chat_completions(request: web.Request) -> web.StreamResponse: - """Handle POST /v1/chat/completions.""" + """Handle POST /v1/chat/completions and /chat/completions.""" body = await request.json() messages = body.get("messages", []) stream = body.get("stream", False) - response_text = match_response(messages) - completion_id = f"mock-{uuid.uuid4().hex[:8]}" + has_tools = bool(body.get("tools")) + cid = f"mock-{uuid.uuid4().hex[:8]}" + # Tool result in messages -> text summary + tr = _find_tool_result(messages) + if tr: + text = f"The {tr['name']} tool returned: {tr['content']}" + if not stream: + return _text_response(cid, text) + return await _stream_text(request, cid, text) + + # Tool-call pattern match + tc = match_tool_call(messages, has_tools) + if tc: + if not stream: + return _tool_call_response(cid, tc) + return await _stream_tool_call(request, cid, tc) + + # Default text response + text = match_response(messages) if not stream: - return web.json_response({ - "id": completion_id, - "object": "chat.completion", - "created": int(time.time()), - "model": "mock-model", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": response_text}, - "finish_reason": "stop", - }], - "usage": {"prompt_tokens": 10, "completion_tokens": len(response_text.split()), "total_tokens": 15}, - }) + return _text_response(cid, text) + return await _stream_text(request, cid, text) - # Streaming response: split into word-boundary chunks - resp = web.StreamResponse( - status=200, - headers={"Content-Type": "text/event-stream", "Cache-Control": "no-cache"}, - ) - await resp.prepare(request) - # First chunk: role - chunk = { - "id": completion_id, - "object": "chat.completion.chunk", - "created": int(time.time()), +def _text_response(cid: str, text: str) -> web.Response: + return web.json_response({ + "id": cid, "object": "chat.completion", "created": int(time.time()), "model": "mock-model", - "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}], - } - await resp.write(f"data: {json.dumps(chunk)}\n\n".encode()) + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": len(text.split()), "total_tokens": 15}, + }) - # Content chunks: split on spaces - words = response_text.split(" ") - for i, word in enumerate(words): - text = word if i == 0 else f" {word}" - chunk["choices"][0]["delta"] = {"content": text} - await resp.write(f"data: {json.dumps(chunk)}\n\n".encode()) - # Final chunk: finish_reason +def _tool_call_response(cid: str, tc: dict) -> web.Response: + return web.json_response({ + "id": cid, "object": "chat.completion", "created": int(time.time()), + "model": "mock-model", + "choices": [{"index": 0, "message": { + "role": "assistant", "content": None, + "tool_calls": [{"id": f"call_{uuid.uuid4().hex[:8]}", "type": "function", + "function": {"name": tc["tool_name"], + "arguments": json.dumps(tc["arguments"])}}], + }, "finish_reason": "tool_calls"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }) + + +async def _stream_text(request: web.Request, cid: str, text: str) -> web.StreamResponse: + resp = web.StreamResponse(status=200, headers={ + "Content-Type": "text/event-stream", "Cache-Control": "no-cache"}) + await resp.prepare(request) + base = _make_base(cid) + chunk = {**base, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, + "finish_reason": None}]} + await _send_sse(resp, chunk) + for i, word in enumerate(text.split(" ")): + chunk["choices"][0]["delta"] = {"content": word if i == 0 else f" {word}"} + await _send_sse(resp, chunk) chunk["choices"][0]["delta"] = {} chunk["choices"][0]["finish_reason"] = "stop" - await resp.write(f"data: {json.dumps(chunk)}\n\n".encode()) + await _send_sse(resp, chunk) await resp.write(b"data: [DONE]\n\n") - return resp +async def _stream_tool_call(request: web.Request, cid: str, tc: dict) -> web.StreamResponse: + resp = web.StreamResponse(status=200, headers={ + "Content-Type": "text/event-stream", "Cache-Control": "no-cache"}) + await resp.prepare(request) + call_id = f"call_{uuid.uuid4().hex[:8]}" + base = _make_base(cid) + # First chunk: role + tool call header with empty arguments + chunk = {**base, "choices": [{"index": 0, "delta": { + "role": "assistant", + "tool_calls": [{"index": 0, "id": call_id, "type": "function", + "function": {"name": tc["tool_name"], "arguments": ""}}], + }, "finish_reason": None}]} + await _send_sse(resp, chunk) + # Second chunk: arguments payload + chunk["choices"][0]["delta"] = { + "tool_calls": [{"index": 0, "function": {"arguments": json.dumps(tc["arguments"])}}]} + await _send_sse(resp, chunk) + # Final chunk: finish reason + chunk["choices"][0]["delta"] = {} + chunk["choices"][0]["finish_reason"] = "tool_calls" + await _send_sse(resp, chunk) + await resp.write(b"data: [DONE]\n\n") + return resp + + +async def oauth_exchange(request: web.Request) -> web.Response: + """Mock OAuth token exchange proxy for E2E tests. + + Accepts form params (code, redirect_uri, code_verifier) and returns + a fake token response. Called by ironclaw's exchange_via_proxy() when + IRONCLAW_OAUTH_EXCHANGE_URL is set. + """ + data = await request.post() + code = data.get("code", "") + return web.json_response({ + "access_token": f"mock-token-{code}", + "refresh_token": "mock-refresh-token", + "expires_in": 3600, + }) + + async def models(_request: web.Request) -> web.Response: - """Handle GET /v1/models.""" return web.json_response({ "object": "list", "data": [{"id": "mock-model", "object": "model", "owned_by": "test"}], @@ -102,23 +229,21 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("--port", type=int, default=0) args = parser.parse_args() - app = web.Application() + # Register both /v1/ and non-/v1/ paths (rig-core omits the /v1/ prefix) app.router.add_post("/v1/chat/completions", chat_completions) + app.router.add_post("/chat/completions", chat_completions) app.router.add_get("/v1/models", models) - - # Use aiohttp's runner to get the actual bound port - import asyncio + app.router.add_get("/models", models) + app.router.add_post("/oauth/exchange", oauth_exchange) async def start(): runner = web.AppRunner(app) await runner.setup() site = web.TCPSite(runner, "127.0.0.1", args.port) await site.start() - # Extract the actual port from the bound socket port = site._server.sockets[0].getsockname()[1] print(f"MOCK_LLM_PORT={port}", flush=True) - # Block forever await asyncio.Event().wait() asyncio.run(start()) diff --git a/tests/e2e/scenarios/test_extension_oauth.py b/tests/e2e/scenarios/test_extension_oauth.py new file mode 100644 index 00000000..b20d4275 --- /dev/null +++ b/tests/e2e/scenarios/test_extension_oauth.py @@ -0,0 +1,264 @@ +"""Extension OAuth round-trip e2e tests. + +Tests the full internal OAuth callback pipeline: install gmail → configure +(get auth_url) → simulate OAuth callback → verify token stored. Uses gateway +callback mode + mock token exchange (no real Google login). + +The conftest sets IRONCLAW_OAUTH_CALLBACK_URL (non-loopback, forces gateway +mode) and IRONCLAW_OAUTH_EXCHANGE_URL (points to mock_llm.py's /oauth/exchange). +""" + +from urllib.parse import parse_qs, urlparse + +import httpx +import pytest + +from helpers import api_get, api_post + +# Module-level state +_gmail_installed = False +_auth_url = None +_csrf_state = None + + +def _extract_state(auth_url: str) -> str: + """Extract the CSRF state parameter from an OAuth authorization URL.""" + parsed = urlparse(auth_url) + qs = parse_qs(parsed.query) + assert "state" in qs, f"auth_url should contain state param: {auth_url}" + state = qs["state"][0] + assert len(state) > 0 + return state + + +async def _get_extension(base_url, name): + """Get a specific extension from the extensions list, or None.""" + r = await api_get(base_url, "/api/extensions") + for ext in r.json().get("extensions", []): + if ext["name"] == name: + return ext + return None + + +async def _ensure_removed(base_url, name): + """Remove extension if already installed.""" + ext = await _get_extension(base_url, name) + if ext: + await api_post(base_url, f"/api/extensions/{name}/remove", timeout=30) + + +# ── Section A: Install + OAuth Initiation ──────────────────────────────── + + +async def test_oauth_install_gmail(ironclaw_server): + """Install gmail from registry for OAuth testing.""" + global _gmail_installed + await _ensure_removed(ironclaw_server, "gmail") + + r = await api_post( + ironclaw_server, + "/api/extensions/install", + json={"name": "gmail"}, + timeout=180, + ) + assert r.status_code == 200 + data = r.json() + assert data.get("success") is True, f"Install failed: {data.get('message', '')}" + _gmail_installed = True + + +async def test_oauth_configure_returns_auth_url(ironclaw_server): + """Configure with empty secrets returns an OAuth auth_url.""" + global _auth_url, _csrf_state + if not _gmail_installed: + pytest.skip("gmail not installed") + + r = await api_post( + ironclaw_server, + "/api/extensions/gmail/setup", + json={"secrets": {}}, + timeout=30, + ) + assert r.status_code == 200 + data = r.json() + assert data.get("success") is True, f"Configure failed: {data.get('message', '')}" + + _auth_url = data.get("auth_url") + assert _auth_url is not None, f"Expected auth_url in response: {data}" + assert "accounts.google.com" in _auth_url, ( + f"auth_url should point to Google: {_auth_url}" + ) + + _csrf_state = _extract_state(_auth_url) + + +async def test_oauth_activate_returns_auth_url(ironclaw_server): + """Activate on un-authenticated gmail returns auth_url.""" + if not _gmail_installed: + pytest.skip("gmail not installed") + + r = await api_post( + ironclaw_server, "/api/extensions/gmail/activate", timeout=30 + ) + assert r.status_code == 200 + data = r.json() + # Activation may fail with auth_url or succeed with auth_url + auth_url = data.get("auth_url") + assert auth_url is not None, f"Expected auth_url in activate response: {data}" + + +# ── Section B: Internal OAuth Round-Trip ───────────────────────────────── + + +async def test_oauth_callback_exchanges_token(ironclaw_server): + """Simulate OAuth callback with mock code — verifies token exchange.""" + global _csrf_state + if not _csrf_state: + pytest.skip("No CSRF state from configure step") + + # Re-configure to get a fresh pending flow (previous configure may have + # been consumed by the activate test above) + r = await api_post( + ironclaw_server, + "/api/extensions/gmail/setup", + json={"secrets": {}}, + timeout=30, + ) + data = r.json() + auth_url = data.get("auth_url") + if auth_url: + _csrf_state = _extract_state(auth_url) + + # Hit the OAuth callback endpoint directly (public route, no auth header). + # The callback handler looks up the pending flow by state, calls + # exchange_via_proxy() which hits mock_llm.py's /oauth/exchange, and + # stores the returned fake token. + async with httpx.AsyncClient() as client: + r = await client.get( + f"{ironclaw_server}/oauth/callback", + params={"code": "mock_auth_code", "state": _csrf_state}, + timeout=30, + follow_redirects=True, + ) + + assert r.status_code == 200, f"Callback returned {r.status_code}: {r.text[:300]}" + body = r.text.lower() + # The landing page says " Connected" on success, "failed" on error + assert "connected" in body or "success" in body, ( + f"Callback HTML should indicate success: {r.text[:500]}" + ) + + +async def test_oauth_callback_replay_rejected(ironclaw_server): + """Replaying the same callback is rejected (flow consumed on first use).""" + if not _csrf_state: + pytest.skip("No CSRF state") + + async with httpx.AsyncClient() as client: + r = await client.get( + f"{ironclaw_server}/oauth/callback", + params={"code": "mock_auth_code", "state": _csrf_state}, + timeout=10, + follow_redirects=True, + ) + + # Should fail — the flow was already consumed + body = r.text.lower() + assert "error" in body or "fail" in body or "expired" in body or r.status_code >= 400, ( + f"Replay should be rejected, got status={r.status_code}: {r.text[:500]}" + ) + + +async def test_oauth_callback_invalid_state(ironclaw_server): + """Callback with bogus state is rejected.""" + async with httpx.AsyncClient() as client: + r = await client.get( + f"{ironclaw_server}/oauth/callback", + params={"code": "x", "state": "totally-bogus-state-value"}, + timeout=10, + follow_redirects=True, + ) + + body = r.text.lower() + assert "error" in body or "fail" in body or "expired" in body or r.status_code >= 400, ( + f"Invalid state should be rejected, got status={r.status_code}: {r.text[:500]}" + ) + + +async def test_oauth_extension_authenticated(ironclaw_server): + """After OAuth callback, gmail shows authenticated=True.""" + if not _gmail_installed: + pytest.skip("gmail not installed") + + ext = await _get_extension(ironclaw_server, "gmail") + assert ext is not None, "gmail not in extensions list" + assert ext["authenticated"] is True, ( + f"gmail should be authenticated after OAuth callback: {ext}" + ) + + +async def test_oauth_tools_registered(ironclaw_server): + """After OAuth authentication, gmail tools appear in tools endpoint.""" + if not _gmail_installed: + pytest.skip("gmail not installed") + + ext = await _get_extension(ironclaw_server, "gmail") + assert ext is not None + # Check the extension's tools array + tools = ext.get("tools", []) + assert len(tools) > 0, ( + f"gmail should have tools registered after auth: {ext}" + ) + + +async def test_remove_during_pending_oauth_invalidates_callback(ironclaw_server): + """Removing an extension while OAuth is pending invalidates the callback state.""" + if not _gmail_installed: + pytest.skip("gmail not installed") + + r = await api_post( + ironclaw_server, + "/api/extensions/gmail/setup", + json={"secrets": {}}, + timeout=30, + ) + assert r.status_code == 200 + data = r.json() + auth_url = data.get("auth_url") + assert auth_url is not None, f"Expected auth_url in response: {data}" + callback_state = _extract_state(auth_url) + + remove_r = await api_post( + ironclaw_server, "/api/extensions/gmail/remove", timeout=30 + ) + assert remove_r.status_code == 200 + assert remove_r.json().get("success") is True, ( + f"Removing gmail during pending OAuth should succeed: {remove_r.text[:300]}" + ) + + async with httpx.AsyncClient() as client: + callback_r = await client.get( + f"{ironclaw_server}/oauth/callback", + params={"code": "mock_auth_code", "state": callback_state}, + timeout=30, + follow_redirects=True, + ) + + assert callback_r.status_code == 200 + body = callback_r.text.lower() + assert "error" in body or "fail" in body or "expired" in body, ( + f"Callback after removal should fail: {callback_r.text[:500]}" + ) + + ext = await _get_extension(ironclaw_server, "gmail") + assert ext is None, "gmail should remain removed after invalidated callback" + + +# ── Section C: Cleanup ────────────────────────────────────────────────── + + +async def test_cleanup_gmail(ironclaw_server): + """Remove gmail (cleanup for other test files).""" + await _ensure_removed(ironclaw_server, "gmail") + ext = await _get_extension(ironclaw_server, "gmail") + assert ext is None, "gmail should be removed" diff --git a/tests/e2e/scenarios/test_extensions.py b/tests/e2e/scenarios/test_extensions.py index 6cddacb4..f172d420 100644 --- a/tests/e2e/scenarios/test_extensions.py +++ b/tests/e2e/scenarios/test_extensions.py @@ -458,6 +458,37 @@ async def test_install_wasm_channel_triggers_configure(page): assert await modal.is_visible() +async def test_install_with_auth_url_opens_popup_and_shows_auth_prompt(page): + """Install responses with auth_url should surface the same auth prompt used elsewhere.""" + await page.evaluate("window.open = (url) => { window._lastOpenedUrl = url; }") + await mock_ext_apis(page, registry=[_REGISTRY_WASM]) + + async def handle_install(route): + await route.fulfill( + status=200, + content_type="application/json", + body=json.dumps({"success": True, "auth_url": "https://example.com/oauth"}), + ) + + await page.route("**/api/extensions/install", handle_install) + await go_to_extensions(page) + + install_btn = page.locator(SEL["available_wasm_list"]).locator(SEL["ext_install_btn"]).first + await install_btn.wait_for(state="visible", timeout=5000) + await install_btn.click() + + await page.wait_for_function( + "() => window._lastOpenedUrl !== null && window._lastOpenedUrl !== undefined", + timeout=5000, + ) + opened = await page.evaluate("window._lastOpenedUrl") + assert opened is not None, "window.open was not called" + assert "example.com" in opened + await page.locator(SEL["auth_card"] + '[data-extension-name="registry-tool"]').wait_for( + state="visible", timeout=5000 + ) + + # ─── Group F: Remove flow ───────────────────────────────────────────────────── async def test_remove_installed_extension_confirmed(page): @@ -612,7 +643,7 @@ async def test_configure_modal_save_success(page): async def test_configure_modal_save_oauth(page): - """Save response with auth_url opens a popup via window.open.""" + """Save response with auth_url opens a popup and shows the global auth prompt.""" await page.evaluate("window.open = (url) => { window._lastOpenedUrl = url; }") async def handle_setup(route): @@ -639,6 +670,9 @@ async def test_configure_modal_save_oauth(page): opened = await page.evaluate("window._lastOpenedUrl") assert opened is not None, "window.open was not called" assert "oauth" in opened or "example.com" in opened + await page.locator(SEL["auth_card"] + '[data-extension-name="test-ext"]').wait_for( + state="visible", timeout=5000 + ) async def test_configure_modal_save_failure(page): @@ -699,7 +733,7 @@ async def test_configure_modal_enter_key_submits(page): # ─── Group H: Auth card (SSE-triggered) ─────────────────────────────────────── async def _show_auth_card(page, **kwargs): - """Inject an auth card via JS and wait for it to appear.""" + """Inject the global auth prompt via JS and wait for it to appear.""" payload = json.dumps(kwargs) await page.evaluate(f"showAuthCard({payload})") await page.locator(SEL["auth_card"]).wait_for(state="visible", timeout=5000) @@ -812,12 +846,43 @@ async def test_auth_card_replaces_existing_same_extension(page): assert "Second" in await page.locator(SEL["auth_instructions"]).text_content() -async def test_auth_card_multiple_extensions_coexist(page): - """Auth cards for different extensions can coexist.""" +async def test_auth_card_for_different_extension_replaces_existing_prompt(page): + """A new auth prompt replaces the previous one to keep the UX modal and global.""" await page.evaluate('showAuthCard({extension_name: "ext-a", instructions: "Token A"})') await page.evaluate('showAuthCard({extension_name: "ext-b", instructions: "Token B"})') - await page.locator(SEL["auth_card"]).nth(1).wait_for(state="visible", timeout=3000) - assert await page.locator(SEL["auth_card"]).count() == 2 + await page.locator(SEL["auth_card"]).wait_for(state="visible", timeout=3000) + assert await page.locator(SEL["auth_card"]).count() == 1 + assert await page.locator(SEL["auth_card"] + '[data-extension-name="ext-a"]').count() == 0 + assert await page.locator(SEL["auth_card"] + '[data-extension-name="ext-b"]').count() == 1 + + +async def test_auth_and_configure_helpers_escape_selector_sensitive_extension_names(page): + """Quoted extension names should not break auth/configure modal helpers.""" + result = await page.evaluate( + """({ name }) => { + showAuthCard({ extension_name: name, instructions: 'Paste token' }); + showAuthCardError(name, 'Bad token'); + const errorText = document.querySelector('.auth-error')?.textContent || ''; + removeAuthCard(name); + const authStillPresent = Array.from(document.querySelectorAll('.auth-card')) + .some((card) => card.getAttribute('data-extension-name') === name); + + const overlay = document.createElement('div'); + overlay.className = 'configure-overlay'; + overlay.setAttribute('data-extension-name', name); + document.body.appendChild(overlay); + closeConfigureModal(name); + const configureStillPresent = Array.from(document.querySelectorAll('.configure-overlay')) + .some((node) => node.getAttribute('data-extension-name') === name); + + return { errorText, authStillPresent, configureStillPresent }; + }""", + {"name": 'quoted "ext" name'}, + ) + + assert result["errorText"] == "Bad token" + assert result["authStillPresent"] is False + assert result["configureStillPresent"] is False async def test_auth_completed_sse_dismisses_card(page): @@ -826,13 +891,95 @@ async def test_auth_completed_sse_dismisses_card(page): # Simulate the auth_completed SSE event being fired await page.evaluate(""" - // Call the handler the same way the SSE listener does - removeAuthCard('myext'); + handleAuthCompleted({ + extension_name: 'myext', + success: true, + message: 'Authenticated!', + }); """) assert await page.locator(SEL["auth_card"] + '[data-extension-name="myext"]').count() == 0 +async def test_auth_completed_for_other_extension_keeps_configure_modal_open(page): + """Auth completion should not close a different extension's configure modal.""" + async def handle_setup(route): + await route.fulfill( + status=200, + content_type="application/json", + body=json.dumps({"secrets": [{"name": "token", "prompt": "Token", "provided": False, "optional": False, "auto_generate": False}]}), + ) + + await page.route("**/api/extensions/test-ext/setup", handle_setup) + await page.evaluate("showConfigureModal('test-ext')") + await page.locator(SEL["configure_modal"]).wait_for(state="visible", timeout=5000) + + await page.evaluate(""" + handleAuthCompleted({ + extension_name: 'other-ext', + success: true, + message: 'Other extension connected.', + }); + """) + + assert await page.locator(SEL["configure_overlay"]).is_visible(), ( + "Configure modal should remain open when another extension finishes auth" + ) + + +async def test_auth_completed_failure_sse_shows_error_toast_and_reloads_extensions(page): + """Failed auth_completed handling should clear stale UI and refresh extensions.""" + reload_count = [] + + async def counting_handler(route): + path = route.request.url.split("?")[0] + if path.endswith("/api/extensions"): + reload_count.append(1) + await route.fulfill( + status=200, + content_type="application/json", + body=json.dumps({"extensions": []}), + ) + else: + await route.continue_() + + async def handle_tools(route): + await route.fulfill(status=200, content_type="application/json", body='{"tools":[]}') + + async def handle_registry(route): + await route.fulfill(status=200, content_type="application/json", body='{"entries":[]}') + + await page.route("**/api/extensions*", counting_handler) + await page.route("**/api/extensions/tools", handle_tools) + await page.route("**/api/extensions/registry", handle_registry) + + await go_to_extensions(page) + count_before = len(reload_count) + + await _show_auth_card(page, extension_name="gmail", auth_url="https://example.com/oauth") + assert await page.locator(SEL["auth_card"] + '[data-extension-name="gmail"]').count() == 1 + + await page.evaluate(""" + handleAuthCompleted({ + extension_name: 'gmail', + success: false, + message: 'OAuth flow expired. Please try again.', + }); + """) + + await wait_for_toast(page, "OAuth flow expired. Please try again.") + assert await page.locator(SEL["auth_card"] + '[data-extension-name="gmail"]').count() == 0 + assert ( + await page.locator( + SEL["toast_error"], has_text="OAuth flow expired. Please try again." + ).count() + >= 1 + ) + + await page.wait_for_timeout(600) + assert len(reload_count) > count_before, "Extensions list did not reload after auth failure" + + # ─── Group I: Activate flow ──────────────────────────────────────────────────── async def test_activate_mcp_server_success(page): @@ -902,8 +1049,8 @@ async def test_activate_failure_shows_error_toast(page): await wait_for_toast(page, "Config missing") -async def test_activate_with_auth_url_opens_popup(page): - """Activate response with auth_url calls window.open.""" +async def test_activate_with_auth_url_opens_popup_and_shows_auth_prompt(page): + """Activate response with auth_url calls window.open and shows the auth prompt.""" await page.evaluate("window.open = (url) => { window._lastOpenedUrl = url; }") await mock_ext_apis(page, installed=[_MCP_INACTIVE]) @@ -921,6 +1068,9 @@ async def test_activate_with_auth_url_opens_popup(page): opened = await page.evaluate("window._lastOpenedUrl") assert opened is not None, "window.open was not called" assert "example.com" in opened + await page.locator( + SEL["auth_card"] + '[data-extension-name="test-mcp-inactive"]' + ).wait_for(state="visible", timeout=5000) # ─── Group J: Tab reload behaviour ──────────────────────────────────────────── @@ -947,9 +1097,9 @@ async def test_extensions_tab_reloads_on_revisit(page): async def handle_registry(route): await route.fulfill(status=200, content_type="application/json", body='{"entries":[]}') + await page.route("**/api/extensions*", counting_handler) await page.route("**/api/extensions/tools", handle_tools) await page.route("**/api/extensions/registry", handle_registry) - await page.route("**/api/extensions*", counting_handler) # First visit await go_to_extensions(page) @@ -990,19 +1140,20 @@ async def test_auth_completed_sse_triggers_extensions_reload(page): async def handle_registry(route): await route.fulfill(status=200, content_type="application/json", body='{"entries":[]}') + await page.route("**/api/extensions*", counting_handler) await page.route("**/api/extensions/tools", handle_tools) await page.route("**/api/extensions/registry", handle_registry) - await page.route("**/api/extensions*", counting_handler) await go_to_extensions(page) count_before = len(reload_count) - # Simulate auth_completed by calling loadExtensions directly (as the SSE handler does) + # Simulate auth_completed via the shared handler. await page.evaluate(""" - // Simulate what the auth_completed SSE handler does when currentTab === 'extensions' - if (typeof loadExtensions === 'function') { - loadExtensions(); - } + handleAuthCompleted({ + extension_name: 'reload-ext', + success: true, + message: 'Reloaded.', + }); """) await page.wait_for_timeout(600) diff --git a/tests/e2e/scenarios/test_pairing.py b/tests/e2e/scenarios/test_pairing.py new file mode 100644 index 00000000..e3ff9144 --- /dev/null +++ b/tests/e2e/scenarios/test_pairing.py @@ -0,0 +1,79 @@ +"""DM pairing flow e2e tests. + +Tests the pairing security gate for WASM channels: listing pending requests, +approving codes, and error handling. +""" + +import httpx +from helpers import AUTH_TOKEN + + +def _headers(): + return {"Authorization": f"Bearer {AUTH_TOKEN}"} + + +async def test_pairing_list_returns_empty_for_unknown_channel(ironclaw_server): + """GET /api/pairing/{channel} returns empty list or 404 for non-existent channel.""" + async with httpx.AsyncClient() as client: + r = await client.get( + f"{ironclaw_server}/api/pairing/nonexistent-channel", + headers=_headers(), + timeout=10, + ) + # Either empty list or error is acceptable + if r.status_code == 200: + data = r.json() + assert isinstance(data, (dict, list)) + if isinstance(data, dict): + assert "requests" in data + assert isinstance(data["requests"], list) + assert data["requests"] == [] + else: + assert data == [] + else: + # 404 or similar is fine for non-existent channel + assert r.status_code in (404, 400) + + +async def test_approve_invalid_code_rejected(ironclaw_server): + """POST /api/pairing/{channel}/approve with bad code returns error.""" + async with httpx.AsyncClient() as client: + r = await client.post( + f"{ironclaw_server}/api/pairing/test-channel/approve", + json={"code": "INVALID0"}, + headers=_headers(), + timeout=10, + ) + # Should fail — no pending request with this code + if r.status_code == 200: + data = r.json() + assert data.get("success") is False or data.get("ok") is False or "error" in str(data).lower() + else: + assert r.status_code >= 400 + + +async def test_approve_empty_code_rejected(ironclaw_server): + """POST /api/pairing/{channel}/approve with empty code returns error.""" + async with httpx.AsyncClient() as client: + r = await client.post( + f"{ironclaw_server}/api/pairing/test-channel/approve", + json={"code": ""}, + headers=_headers(), + timeout=10, + ) + if r.status_code == 200: + data = r.json() + assert data.get("success") is False or data.get("ok") is False + else: + assert r.status_code >= 400 + + +async def test_pairing_approve_requires_auth(ironclaw_server): + """POST /api/pairing/{channel}/approve without auth token is rejected.""" + async with httpx.AsyncClient() as client: + r = await client.post( + f"{ironclaw_server}/api/pairing/test-channel/approve", + json={"code": "ABCD1234"}, + timeout=10, + ) + assert r.status_code == 401 or r.status_code == 403 diff --git a/tests/e2e/scenarios/test_tool_execution.py b/tests/e2e/scenarios/test_tool_execution.py new file mode 100644 index 00000000..89627ac3 --- /dev/null +++ b/tests/e2e/scenarios/test_tool_execution.py @@ -0,0 +1,94 @@ +"""Tool execution e2e tests. + +Tests the agent loop: user message -> mock LLM returns tool_calls -> tool +executes -> result displayed in chat. Requires the enhanced mock_llm.py +with TOOL_CALL_PATTERNS support. +""" + +from helpers import SEL + + +async def _send_and_get_response( + page, + message: str, + *, + expected_fragment: str, + timeout: int = 30000, +) -> str: + """Send a message and return the text of the newest assistant response. + + Counts existing assistant messages before sending, then waits for a new + one to appear and contain the expected final text fragment. This avoids + reading partial streamed content before the assistant response is complete. + """ + chat_input = page.locator(SEL["chat_input"]) + await chat_input.wait_for(state="visible", timeout=5000) + + # Count existing assistant messages before sending + assistant_sel = SEL["message_assistant"] + before_count = await page.locator(assistant_sel).count() + + await chat_input.fill(message) + await chat_input.press("Enter") + + # Wait for the final assistant message to exist and include the expected + # text fragment rather than returning on the first streamed chunk. + expected = before_count + 1 + await page.wait_for_function( + """({ assistantSelector, expectedCount, expectedFragment }) => { + const messages = document.querySelectorAll(assistantSelector); + if (messages.length < expectedCount) return false; + const text = (messages[messages.length - 1].innerText || '').trim().toLowerCase(); + return text.includes(expectedFragment.toLowerCase()); + }""", + arg={ + "assistantSelector": assistant_sel, + "expectedCount": expected, + "expectedFragment": expected_fragment, + }, + timeout=timeout, + ) + + return await page.locator(assistant_sel).last.inner_text() + + +async def test_builtin_echo_tool(page): + """Send a message that triggers the echo tool via mock LLM function calling.""" + text = await _send_and_get_response( + page, + "echo hello world", + expected_fragment="hello world", + ) + + # The mock LLM returns "The echo tool returned: " + assert "echo" in text.lower() or "hello world" in text.lower(), ( + f"Expected echo result in response, got: {text}" + ) + + +async def test_builtin_time_tool(page): + """Send a message that triggers the time tool via mock LLM function calling.""" + text = await _send_and_get_response( + page, + "what time is it", + expected_fragment="time", + ) + + # The mock LLM returns "The time tool returned: " + assert "time" in text.lower(), ( + f"Expected time result in response, got: {text}" + ) + + +async def test_non_tool_message_still_works(page): + """Messages that don't match tool patterns still get text responses.""" + text = await _send_and_get_response( + page, + "What is 2+2?", + expected_fragment="4", + timeout=15000, + ) + + assert "4" in text, ( + f"Expected '4' in response, got: {text}" + ) diff --git a/tests/e2e/scenarios/test_wasm_lifecycle.py b/tests/e2e/scenarios/test_wasm_lifecycle.py new file mode 100644 index 00000000..961e7ad0 --- /dev/null +++ b/tests/e2e/scenarios/test_wasm_lifecycle.py @@ -0,0 +1,517 @@ +"""Comprehensive WASM extension lifecycle e2e tests. + +Tests the full extension pipeline: registry → install → fields → configure → +activate → tools → remove → reinstall. Validates response fields, not just +status codes, to catch production bugs like missing capabilities, wrong +activation state, and stale registry flags. + +Lifecycle stages are expressed as scoped fixtures so each test requests the +state it needs explicitly rather than relying on module-global flags. +""" + +from pathlib import Path + +import pytest + +from helpers import SEL, api_get, api_post + +async def _get_extension(base_url, name): + """Get a specific extension from the extensions list, or None.""" + r = await api_get(base_url, "/api/extensions") + for ext in r.json().get("extensions", []): + if ext["name"] == name: + return ext + return None + + +async def _ensure_removed(base_url, name): + """Remove extension if already installed (idempotent cleanup).""" + ext = await _get_extension(base_url, name) + if ext: + await api_post(base_url, f"/api/extensions/{name}/remove", timeout=30) + + +async def _install_extension(base_url, name): + """Install an extension and assert success.""" + r = await api_post( + base_url, + "/api/extensions/install", + json={"name": name}, + timeout=180, + ) + assert r.status_code == 200, f"Install HTTP error: {r.status_code} {r.text[:300]}" + data = r.json() + assert data.get("success") is True, f"Install failed: {data.get('message', '')}" + return data + + +@pytest.fixture(scope="module", autouse=True) +async def extension_lifecycle_cleanup(ironclaw_server): + """Start and end the module with a clean extension set.""" + await _ensure_removed(ironclaw_server, "web-search") + await _ensure_removed(ironclaw_server, "gmail") + yield + await _ensure_removed(ironclaw_server, "web-search") + await _ensure_removed(ironclaw_server, "gmail") + + +@pytest.fixture(scope="module") +async def web_search_installed(ironclaw_server, extension_lifecycle_cleanup): + """Install web-search once for tests that require the pre-configure state.""" + data = await _install_extension(ironclaw_server, "web-search") + return {"name": "web-search", "install": data} + + +@pytest.fixture(scope="module") +async def web_search_configured(ironclaw_server, web_search_installed): + """Configure web-search once for tests that require the active state.""" + r = await api_post( + ironclaw_server, + "/api/extensions/web-search/setup", + json={"secrets": {"brave_api_key": "test-key-123"}}, + timeout=30, + ) + assert r.status_code == 200 + data = r.json() + assert data.get("success") is True, f"Configure failed: {data.get('message', '')}" + assert data.get("activated") is True, "Should auto-activate after configure" + return {"name": "web-search", "configure": data} + + +@pytest.fixture(scope="module") +async def gmail_installed(ironclaw_server, extension_lifecycle_cleanup): + """Install gmail once for multi-extension and OAuth setup assertions.""" + data = await _install_extension(ironclaw_server, "gmail") + return {"name": "gmail", "install": data} + + +@pytest.fixture(scope="module") +async def web_search_removed(ironclaw_server, web_search_configured): + """Remove web-search once for post-uninstall assertions.""" + r = await api_post( + ironclaw_server, "/api/extensions/web-search/remove", timeout=30 + ) + assert r.status_code == 200 + data = r.json() + assert data.get("success") is True, f"Remove failed: {data.get('message', '')}" + return {"name": "web-search", "remove": data} + + +@pytest.fixture(scope="module") +async def web_search_reinstalled(ironclaw_server, web_search_removed): + """Reinstall web-search after removal to verify saved-secret recovery.""" + await _ensure_removed(ironclaw_server, "web-search") + data = await _install_extension(ironclaw_server, "web-search") + return {"name": "web-search", "install": data} + + +# ── Section A: Registry Validation ────────────────────────────────────── + + +async def test_registry_lists_extensions(ironclaw_server): + """Registry endpoint returns entries from the embedded catalog.""" + r = await api_get(ironclaw_server, "/api/extensions/registry") + assert r.status_code == 200 + data = r.json() + assert "entries" in data + names = [e["name"] for e in data["entries"]] + assert "web-search" in names + assert "gmail" in names + + +async def test_registry_entry_fields(ironclaw_server): + """Every registry entry has all required fields with correct types.""" + r = await api_get(ironclaw_server, "/api/extensions/registry") + entries = r.json()["entries"] + assert len(entries) > 0, "Registry should have entries" + for entry in entries: + assert "name" in entry and isinstance(entry["name"], str) and entry["name"] + assert "display_name" in entry and isinstance(entry["display_name"], str) + assert "kind" in entry and isinstance(entry["kind"], str) + assert "description" in entry and isinstance(entry["description"], str) + assert "installed" in entry and isinstance(entry["installed"], bool) + assert "keywords" in entry and isinstance(entry["keywords"], list) + + +async def test_registry_installed_flag_false_initially(ironclaw_server): + """Before any install, all registry entries have installed=False.""" + # Clean up in case previous test run left extensions installed + await _ensure_removed(ironclaw_server, "web-search") + await _ensure_removed(ironclaw_server, "gmail") + + r = await api_get(ironclaw_server, "/api/extensions/registry") + entries = r.json()["entries"] + for entry in entries: + if entry["name"] in ("web-search", "gmail"): + assert entry["installed"] is False, ( + f"{entry['name']} should not be installed yet" + ) + + +async def test_registry_search_filters(ironclaw_server): + """Search query filters registry results.""" + r = await api_get( + ironclaw_server, "/api/extensions/registry", params={"query": "search"} + ) + assert r.status_code == 200 + entries = r.json()["entries"] + names = [e["name"] for e in entries] + assert "web-search" in names + + +async def test_registry_search_no_match(ironclaw_server): + """Nonsense query returns empty results.""" + r = await api_get( + ironclaw_server, + "/api/extensions/registry", + params={"query": "xyznonexistent999"}, + ) + assert r.status_code == 200 + assert len(r.json()["entries"]) == 0 + + +# ── Section B: Install Lifecycle (web-search) ─────────────────────────── + + +async def test_install_web_search(web_search_installed): + """Install web-search from registry. Asserts success — failure here means + the registry/download/build pipeline is broken.""" + assert "message" in web_search_installed["install"] + + +async def test_installed_extension_fields(ironclaw_server, web_search_installed): + """After install, extension list shows correct fields.""" + ext = await _get_extension(ironclaw_server, "web-search") + assert ext is not None, "web-search not in extensions list after install" + assert ext["kind"] == "wasm_tool" + assert ext["needs_setup"] is True, "Should need setup (has brave_api_key secret)" + assert ext["authenticated"] is False, "Should not be authenticated before configure" + + +async def test_installed_in_registry(ironclaw_server, web_search_installed): + """Registry marks installed extension with installed=True.""" + r = await api_get(ironclaw_server, "/api/extensions/registry") + entries = r.json()["entries"] + ws_entry = next((e for e in entries if e["name"] == "web-search"), None) + assert ws_entry is not None + assert ws_entry["installed"] is True, "Registry should show installed=True" + + +async def test_setup_schema_has_secrets(ironclaw_server, web_search_installed): + """Setup schema returns brave_api_key with correct field info.""" + r = await api_get(ironclaw_server, "/api/extensions/web-search/setup") + assert r.status_code == 200 + data = r.json() + assert "secrets" in data + secrets = {s["name"]: s for s in data["secrets"]} + assert "brave_api_key" in secrets, ( + f"brave_api_key not in setup schema secrets: {list(secrets.keys())}" + ) + key_info = secrets["brave_api_key"] + assert key_info["provided"] is False, "Should not be provided yet" + + +async def test_extension_not_authenticated_before_configure( + ironclaw_server, web_search_installed +): + """Installed but not configured extension is not authenticated.""" + ext = await _get_extension(ironclaw_server, "web-search") + assert ext is not None + # Before configuring secrets, extension shouldn't be fully authenticated + assert ext["needs_setup"] is True, "Should still need setup before configure" + + +async def test_activate_before_configure_rejected(ironclaw_server, web_search_installed): + """Activating a tool that needs setup secrets is rejected.""" + r = await api_post( + ironclaw_server, "/api/extensions/web-search/activate", timeout=30 + ) + assert r.status_code == 200 + data = r.json() + assert data.get("success") is False, ( + f"Activate should fail before configure: {data}" + ) + msg = data.get("message", "").lower() + assert "requires configuration" in msg or "setup" in msg, ( + f"Error should mention configuration: {data.get('message')}" + ) + + +# ── Section C: Configure + Activate (web-search) ──────────────────────── + + +async def test_configure_rejects_unknown_secret(ironclaw_server, web_search_installed): + """Submitting an unknown secret name is rejected.""" + r = await api_post( + ironclaw_server, + "/api/extensions/web-search/setup", + json={"secrets": {"fake_unknown_key": "value"}}, + ) + assert r.status_code == 200 + data = r.json() + assert data.get("success") is False, f"Should reject unknown secret: {data}" + assert "unknown" in data.get("message", "").lower() or "not found" in data.get( + "message", "" + ).lower(), f"Error should mention unknown secret: {data.get('message')}" + + +async def test_configure_with_valid_secret(web_search_configured): + """Configure with valid brave_api_key succeeds and auto-activates.""" + assert web_search_configured["configure"].get("activated") is True + + +async def test_extension_active_after_configure(ironclaw_server, web_search_configured): + """After configure, extension shows authenticated=True and active=True.""" + ext = await _get_extension(ironclaw_server, "web-search") + assert ext is not None + assert ext["authenticated"] is True, "Should be authenticated after configure" + assert ext["active"] is True, "Should be active after auto-activation" + assert len(ext.get("tools", [])) > 0, "Should have tools registered" + + +async def test_setup_shows_provided(ironclaw_server, web_search_configured): + """After configure, setup schema shows secret as provided.""" + r = await api_get(ironclaw_server, "/api/extensions/web-search/setup") + assert r.status_code == 200 + secrets = {s["name"]: s for s in r.json()["secrets"]} + assert "brave_api_key" in secrets + assert secrets["brave_api_key"]["provided"] is True + + +async def test_tools_registered_after_activate( + ironclaw_server, web_search_configured +): + """After activation, extension tools appear in the tools endpoint.""" + r = await api_get(ironclaw_server, "/api/extensions/tools") + assert r.status_code == 200 + tool_names = [t["name"] for t in r.json()["tools"]] + assert "web-search" in tool_names, ( + f"web-search tool not found in tools list: {tool_names}" + ) + + +async def test_activate_already_active_idempotent( + ironclaw_server, web_search_configured +): + """Activating an already-active extension succeeds (idempotent).""" + r = await api_post( + ironclaw_server, "/api/extensions/web-search/activate", timeout=30 + ) + assert r.status_code == 200 + data = r.json() + assert data.get("success") is True, ( + f"Re-activation should succeed: {data.get('message', '')}" + ) + + +async def test_configure_empty_secret_skipped(ironclaw_server, web_search_configured): + """Submitting an empty string for a secret skips it (doesn't overwrite).""" + r = await api_post( + ironclaw_server, + "/api/extensions/web-search/setup", + json={"secrets": {"brave_api_key": ""}}, + timeout=30, + ) + assert r.status_code == 200 + data = r.json() + assert data.get("success") is True + + # Verify the secret is still provided (not cleared) + r2 = await api_get(ironclaw_server, "/api/extensions/web-search/setup") + secrets = {s["name"]: s for s in r2.json()["secrets"]} + assert secrets["brave_api_key"]["provided"] is True, ( + "Empty value should not clear existing secret" + ) + + +# ── Section D: Install gmail (multi-extension) ────────────────────────── + + +async def test_install_gmail(gmail_installed): + """Install gmail from registry (second extension, tests isolation).""" + assert "message" in gmail_installed["install"] + + +async def test_gmail_fields(ironclaw_server, gmail_installed): + """Gmail extension has correct field values (OAuth-based auth).""" + ext = await _get_extension(ironclaw_server, "gmail") + assert ext is not None, "gmail not in extensions list" + assert ext["kind"] == "wasm_tool" + assert ext["has_auth"] is True, "Gmail should have OAuth auth" + + +async def test_both_extensions_listed( + ironclaw_server, web_search_configured, gmail_installed +): + """Both web-search and gmail appear in extensions list (no clobbering).""" + r = await api_get(ironclaw_server, "/api/extensions") + names = [e["name"] for e in r.json()["extensions"]] + assert "web-search" in names, f"web-search missing from: {names}" + assert "gmail" in names, f"gmail missing from: {names}" + + +async def test_gmail_setup_schema_auto_resolves(ironclaw_server, gmail_installed): + """Gmail setup schema returns empty secrets (builtin creds auto-resolve).""" + r = await api_get(ironclaw_server, "/api/extensions/gmail/setup") + assert r.status_code == 200 + data = r.json() + secrets = data.get("secrets", []) + # Builtin Google credentials auto-resolve client_id/client_secret via + # is_auto_resolved_oauth_field(), so the setup schema should have no + # user-facing secrets (or only auto-generated ones). + user_facing = [s for s in secrets if not s.get("auto_generate", False)] + assert len(user_facing) == 0, ( + f"Gmail should have no user-facing secrets (auto-resolved), got: " + f"{[s['name'] for s in user_facing]}" + ) + + +# ── Section E: Remove + Cleanup ───────────────────────────────────────── + + +async def test_remove_web_search(web_search_removed): + """Remove web-search succeeds.""" + assert web_search_removed["remove"].get("success") is True + + +async def test_removed_not_in_extensions(ironclaw_server, web_search_removed): + """Removed extension no longer appears in extensions list.""" + ext = await _get_extension(ironclaw_server, "web-search") + assert ext is None, "web-search should not be in extensions list after removal" + + +async def test_removed_extension_not_listed(ironclaw_server, web_search_removed): + """Removed extension should not appear in the extension tools list.""" + r = await api_get(ironclaw_server, "/api/extensions/tools") + assert r.status_code == 200 + tool_names = [t["name"] for t in r.json()["tools"]] + assert "web-search" not in tool_names, ( + f"Removed web-search tool should not remain registered: {tool_names}" + ) + + +async def test_removed_not_in_registry_installed(ironclaw_server, web_search_removed): + """Registry shows removed extension as installed=False.""" + r = await api_get(ironclaw_server, "/api/extensions/registry") + ws_entry = next( + (e for e in r.json()["entries"] if e["name"] == "web-search"), None + ) + assert ws_entry is not None + assert ws_entry["installed"] is False, "Registry should show installed=False" + + +async def test_activate_after_remove_uses_replacement_bytes_not_cached_module( + ironclaw_server, wasm_tools_dir, web_search_removed +): + """After removal, activation must use the replacement bytes rather than a stale cache.""" + wasm_path = Path(wasm_tools_dir) / "web-search.wasm" + wasm_path.write_bytes(b"not-a-valid-wasm-component") + + r = await api_post( + ironclaw_server, "/api/extensions/web-search/activate", timeout=30 + ) + assert r.status_code == 200 + data = r.json() + assert data.get("success") is False, ( + f"Activation should fail against replacement bytes, got: {data}" + ) + + +async def test_reinstall_after_remove(ironclaw_server, web_search_reinstalled): + """Extension can be reinstalled after removal without stale activation errors.""" + ext = await _get_extension(ironclaw_server, "web-search") + assert ext is not None, "web-search not found after reinstall" + assert ext["active"] is True, "Reinstalled tool should auto-activate via saved secrets" + assert ext["authenticated"] is True, "Saved secret should still authenticate on reinstall" + # Verify no stale activation error from previous install + assert ext.get("activation_error") is None or ext.get("activation_error") == "", ( + f"Reinstalled extension should have no stale activation error: {ext}" + ) + + +# ── Section F: Error Paths ────────────────────────────────────────────── + + +async def test_install_nonexistent(ironclaw_server): + """Installing a nonexistent extension returns an error.""" + r = await api_post( + ironclaw_server, + "/api/extensions/install", + json={"name": "nonexistent-tool-xyz-999"}, + timeout=30, + ) + if r.status_code == 200: + assert r.json().get("success") is False + else: + assert r.status_code >= 400 + + +async def test_install_empty_name(ironclaw_server): + """Installing with empty name returns an error.""" + r = await api_post( + ironclaw_server, + "/api/extensions/install", + json={"name": ""}, + timeout=10, + ) + if r.status_code == 200: + assert r.json().get("success") is False + else: + assert r.status_code >= 400 + + +async def test_remove_noninstalled(ironclaw_server): + """Removing a non-installed extension returns an error.""" + r = await api_post( + ironclaw_server, "/api/extensions/nonexistent-xyz/remove", timeout=10 + ) + if r.status_code == 200: + assert r.json().get("success") is False + else: + assert r.status_code >= 400 + + +async def test_activate_noninstalled(ironclaw_server): + """Activating a non-installed extension returns an error.""" + r = await api_post( + ironclaw_server, "/api/extensions/nonexistent-xyz/activate", timeout=10 + ) + if r.status_code == 200: + assert r.json().get("success") is False + else: + assert r.status_code >= 400 + + +async def test_setup_noninstalled(ironclaw_server): + """Setup for non-installed extension returns an error.""" + r = await api_get(ironclaw_server, "/api/extensions/nonexistent-xyz/setup") + # May return 500 or a JSON error + assert r.status_code >= 400 or r.json().get("success") is False + + +async def test_configure_noninstalled(ironclaw_server): + """Configure for non-installed extension returns an error.""" + r = await api_post( + ironclaw_server, + "/api/extensions/nonexistent-xyz/setup", + json={"secrets": {}}, + timeout=10, + ) + if r.status_code == 200: + assert r.json().get("success") is False + else: + assert r.status_code >= 400 + + +# ── Section G: Browser UI ────────────────────────────────────────────── + + +async def test_extensions_tab_shows_registry(page): + """Extensions tab loads and shows available extensions from registry.""" + tab_btn = page.locator(SEL["tab_button"].format(tab="extensions")) + await tab_btn.click() + panel = page.locator(SEL["tab_panel"].format(tab="extensions")) + await panel.wait_for(state="visible", timeout=5000) + + available_section = page.locator(SEL["available_wasm_list"]) + await available_section.wait_for(state="visible", timeout=10000) diff --git a/tests/e2e_builtin_tool_coverage.rs b/tests/e2e_builtin_tool_coverage.rs index 34cb35f7..f1ae3660 100644 --- a/tests/e2e_builtin_tool_coverage.rs +++ b/tests/e2e_builtin_tool_coverage.rs @@ -457,4 +457,90 @@ mod tests { rig.shutdown(); } + + // ----------------------------------------------------------------------- + // Test: tool_info_discovery (two-level detail) + // ----------------------------------------------------------------------- + // Verifies the tool_info built-in returns: + // - Default (no include_schema): name, description, parameter names array + // - With include_schema: true: adds full typed JSON Schema + + #[tokio::test] + async fn tool_info_discovery() { + let trace = LlmTrace::from_file(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/llm_traces/tools/tool_info_discovery.json" + )) + .expect("failed to load tool_info_discovery.json"); + + let rig = TestRigBuilder::new() + .with_trace(trace.clone()) + .with_auto_approve_tools(true) + .build() + .await; + + rig.send_message("What is the schema for the echo and time tools?") + .await; + let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await; + + rig.verify_trace_expects(&trace, &responses); + + // tool_info should have been called twice (echo + time), both succeeding. + let completed = rig.tool_calls_completed(); + let tool_info_calls: Vec<_> = completed.iter().filter(|(n, _)| n == "tool_info").collect(); + assert_eq!( + tool_info_calls.len(), + 2, + "Expected 2 tool_info calls, got {tool_info_calls:?}" + ); + assert!( + tool_info_calls.iter().all(|(_, ok)| *ok), + "All tool_info calls should succeed: {tool_info_calls:?}" + ); + + // Verify the results contain expected fields. + let results = rig.tool_results(); + let info_results: Vec<_> = results.iter().filter(|(n, _)| n == "tool_info").collect(); + + // First call was for "echo" (default, no include_schema) — result should + // contain "echo" and "parameters" as an array of names (not full schema). + let echo_result = info_results + .iter() + .find(|(_, preview)| preview.contains("echo")) + .expect("tool_info result should contain 'echo'"); + assert!( + echo_result.1.contains("message"), + "echo default result should list 'message' parameter name: {:?}", + echo_result.1 + ); + // Default mode should NOT include the full "schema" key + let echo_json: serde_json::Value = serde_json::from_str(&echo_result.1) + .expect("echo tool_info result should be valid JSON"); + assert!( + echo_json.get("schema").is_none(), + "Default tool_info should not include schema field: {:?}", + echo_result.1 + ); + + // Second call was for "time" with include_schema: true — result should + // contain "time", "schema" field with full object. + let time_result = info_results + .iter() + .find(|(_, preview)| preview.contains("time")) + .expect("tool_info result should contain 'time'"); + let time_json: serde_json::Value = serde_json::from_str(&time_result.1) + .expect("time tool_info result should be valid JSON"); + assert!( + time_json.get("schema").is_some(), + "include_schema: true should include schema field: {:?}", + time_result.1 + ); + assert!( + time_json["schema"]["properties"].is_object(), + "schema should have properties: {:?}", + time_result.1 + ); + + rig.shutdown(); + } } diff --git a/tests/e2e_routine_heartbeat.rs b/tests/e2e_routine_heartbeat.rs index f245e656..f5a28c25 100644 --- a/tests/e2e_routine_heartbeat.rs +++ b/tests/e2e_routine_heartbeat.rs @@ -61,6 +61,8 @@ mod tests { prompt: prompt.to_string(), context_paths: vec![], max_tokens: 1000, + use_tools: false, + max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: Duration::from_secs(0), diff --git a/tests/fixtures/llm_traces/tools/tool_info_discovery.json b/tests/fixtures/llm_traces/tools/tool_info_discovery.json new file mode 100644 index 00000000..dc8746ad --- /dev/null +++ b/tests/fixtures/llm_traces/tools/tool_info_discovery.json @@ -0,0 +1,50 @@ +{ + "model_name": "test-tool-info-discovery", + "expects": { + "tools_used": ["tool_info"], + "all_tools_succeeded": true, + "min_responses": 1, + "tool_results_contain": { + "tool_info": "echo" + } + }, + "steps": [ + { + "request_hint": { "last_user_message_contains": "schema" }, + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_tool_info_echo", + "name": "tool_info", + "arguments": { "name": "echo" } + } + ], + "input_tokens": 100, + "output_tokens": 20 + } + }, + { + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_tool_info_time", + "name": "tool_info", + "arguments": { "name": "time", "include_schema": true } + } + ], + "input_tokens": 200, + "output_tokens": 20 + } + }, + { + "response": { + "type": "text", + "content": "I found the info for both tools. The echo tool has a 'message' parameter. The time tool accepts an 'operation' parameter with options like 'now', 'parse', and 'diff'.", + "input_tokens": 400, + "output_tokens": 40 + } + } + ] +} diff --git a/tests/gateway_workflow_integration.rs b/tests/gateway_workflow_integration.rs index 2f1353a5..187cc751 100644 --- a/tests/gateway_workflow_integration.rs +++ b/tests/gateway_workflow_integration.rs @@ -13,6 +13,8 @@ mod support; mod tests { use std::time::Duration; + use uuid::Uuid; + use crate::support::gateway_workflow_harness::GatewayWorkflowHarness; use crate::support::mock_openai_server::{ MockOpenAiResponse, MockOpenAiRule, MockOpenAiServerBuilder, MockToolCall, @@ -147,4 +149,115 @@ mod tests { harness.shutdown().await; mock.shutdown().await; } + + #[tokio::test] + async fn routines_toggle_reenable_cron_recomputes_next_fire_at() { + let mock = MockOpenAiServerBuilder::new() + .with_rule(MockOpenAiRule::on_user_contains( + "create cron routine", + MockOpenAiResponse::ToolCalls(vec![MockToolCall::new( + "call_create_cron_1", + "routine_create", + serde_json::json!({ + "name": "wf-cron-toggle-reenable", + "description": "Cron toggle regression test", + "trigger_type": "cron", + "schedule": "0 */5 * * * *", + "timezone": "UTC", + "action_type": "lightweight", + "prompt": "noop" + }), + )]), + )) + .with_default_response(MockOpenAiResponse::Text("ack".to_string())) + .start() + .await; + + let harness = + GatewayWorkflowHarness::start_openai_compatible(&mock.openai_base_url(), "mock-model") + .await; + + let thread_id = harness.create_thread().await; + harness.send_chat(&thread_id, "create cron routine").await; + harness + .wait_for_turns(&thread_id, 1, Duration::from_secs(10)) + .await; + + let routine = harness + .routine_by_name("wf-cron-toggle-reenable") + .await + .expect("routine should exist"); + let routine_id = routine + .get("id") + .and_then(|v| v.as_str()) + .expect("routine id missing"); + + let routine_uuid = Uuid::parse_str(routine_id).expect("valid routine uuid"); + + // Disable through the web toggle endpoint. + harness + .client + .post(format!( + "{}/api/routines/{routine_id}/toggle", + harness.base_url() + )) + .bearer_auth(&harness.auth_token) + .json(&serde_json::json!({ "enabled": false })) + .send() + .await + .expect("disable toggle request failed") + .error_for_status() + .expect("disable toggle non-2xx"); + + // Simulate an unscheduled disabled cron routine (next_fire_at missing). + let mut stored = harness + .db + .get_routine(routine_uuid) + .await + .expect("db get_routine") + .expect("routine should still exist"); + stored.next_fire_at = None; + harness + .db + .update_routine(&stored) + .await + .expect("db update_routine"); + + // Re-enable through the web toggle endpoint. + harness + .client + .post(format!( + "{}/api/routines/{routine_id}/toggle", + harness.base_url() + )) + .bearer_auth(&harness.auth_token) + .json(&serde_json::json!({ "enabled": true })) + .send() + .await + .expect("enable toggle request failed") + .error_for_status() + .expect("enable toggle non-2xx"); + + let detail = harness + .client + .get(format!("{}/api/routines/{routine_id}", harness.base_url())) + .bearer_auth(&harness.auth_token) + .send() + .await + .expect("detail request failed") + .error_for_status() + .expect("detail non-2xx") + .json::() + .await + .expect("invalid detail response"); + + assert_eq!(detail["enabled"].as_bool(), Some(true)); + assert!( + detail["next_fire_at"].as_str().is_some(), + "expected next_fire_at to be recomputed when re-enabling cron routine, got {detail}" + ); + + harness.shutdown().await; + mock.shutdown().await; + } } diff --git a/tools-src/llm-context/Cargo.toml b/tools-src/llm-context/Cargo.toml new file mode 100644 index 00000000..9f672189 --- /dev/null +++ b/tools-src/llm-context/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "llm-context-tool" +version = "0.1.0" +edition = "2021" +description = "Brave Search LLM Context tool for IronClaw (WASM component)" +license = "MIT OR Apache-2.0" +publish = false + +[dependencies] +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +wit-bindgen = "0.41.0" + +[lib] +crate-type = ["cdylib"] + +[profile.release] +opt-level = "s" +lto = true +strip = true +codegen-units = 1 + +[workspace] diff --git a/tools-src/llm-context/llm-context-tool.capabilities.json b/tools-src/llm-context/llm-context-tool.capabilities.json new file mode 100644 index 00000000..72061eaa --- /dev/null +++ b/tools-src/llm-context/llm-context-tool.capabilities.json @@ -0,0 +1,53 @@ +{ + "version": "0.1.0", + "wit_version": "0.3.0", + "capabilities": { + "http": { + "allowlist": [ + { + "host": "api.search.brave.com", + "path_prefix": "/res/v1/llm/context", + "methods": [ + "POST" + ] + } + ], + "credentials": { + "brave_api_key": { + "secret_name": "brave_api_key", + "location": { + "type": "header", + "name": "X-Subscription-Token" + }, + "host_patterns": [ + "api.search.brave.com" + ] + } + }, + "rate_limit": { + "requests_per_minute": 30, + "requests_per_hour": 500 + } + }, + "secrets": { + "allowed_names": [ + "brave_api_key" + ] + } + }, + "auth": { + "secret_name": "brave_api_key", + "display_name": "Brave Search", + "instructions": "Get a free API key at brave.com/search/api/ (Free tier: 2,000 queries/month). Same key as Web Search.", + "setup_url": "https://brave.com/search/api/", + "env_var": "BRAVE_API_KEY" + }, + "setup": { + "required_secrets": [ + { + "name": "brave_api_key", + "prompt": "Brave Search API key (from brave.com/search/api)" + } + ] + } +} diff --git a/tools-src/llm-context/src/lib.rs b/tools-src/llm-context/src/lib.rs new file mode 100644 index 00000000..59791f3b --- /dev/null +++ b/tools-src/llm-context/src/lib.rs @@ -0,0 +1,1339 @@ +//! Brave Search LLM Context WASM Tool for IronClaw. +//! +//! Fetches pre-extracted web content from the Brave Search LLM Context API, +//! optimized for grounding LLM responses (RAG, fact-checking, research). +//! +//! # Authentication +//! +//! Uses the same Brave Search API key as the Web Search tool: +//! `ironclaw secret set brave_api_key ` +//! +//! Get a key at: https://brave.com/search/api/ + +wit_bindgen::generate!({ + world: "sandboxed-tool", + path: "../../wit/tool.wit", +}); + +use serde::Deserialize; + +// Brave LLM Context API endpoint documentation: +// https://api-dashboard.search.brave.com/documentation/services/llm-context +// +// This tool uses POST with a JSON body (unlike Web Search's GET + query params) to avoid +// URL length limits and support richer parameters. + +const BRAVE_LLM_CONTEXT_ENDPOINT: &str = "https://api.search.brave.com/res/v1/llm/context"; + +// Query and result limits (aligned with Brave API) +const MAX_QUERY_LEN: usize = 400; +const MAX_QUERY_WORDS: usize = 50; +const MIN_COUNT: u32 = 1; +const MAX_COUNT: u32 = 50; +const DEFAULT_COUNT: u32 = 20; +const MIN_TOKENS: u32 = 1024; +const MAX_TOKENS: u32 = 32768; +const DEFAULT_MAX_TOKENS: u32 = 8192; +const MIN_URLS: u32 = 1; +const MAX_URLS: u32 = 50; +const DEFAULT_MAX_URLS: u32 = 20; +const MIN_SNIPPETS: u32 = 1; +const MAX_SNIPPETS: u32 = 100; +const DEFAULT_MAX_SNIPPETS: u32 = 50; +const MIN_TOKENS_PER_URL: u32 = 512; +const MAX_TOKENS_PER_URL: u32 = 8192; +const DEFAULT_MAX_TOKENS_PER_URL: u32 = 4096; +const MIN_SNIPPETS_PER_URL: u32 = 1; +const MAX_SNIPPETS_PER_URL: u32 = 100; +const DEFAULT_SNIPPETS_PER_URL: u32 = 50; +const MAX_RETRIES: u32 = 3; + +// Validation helpers +const VALID_THRESHOLD_MODES: [&str; 4] = ["strict", "balanced", "lenient", "disabled"]; + +struct LlmContextTool; + +impl exports::near::agent::tool::Guest for LlmContextTool { + fn execute(req: exports::near::agent::tool::Request) -> exports::near::agent::tool::Response { + match execute_inner(&req.params) { + Ok(result) => exports::near::agent::tool::Response { + output: Some(result), + error: None, + }, + Err(e) => exports::near::agent::tool::Response { + output: None, + error: Some(e), + }, + } + } + + fn schema() -> String { + SCHEMA.to_string() + } + + fn description() -> String { + "Fetch pre-extracted web content from Brave Search for grounding LLM answers. \ + Returns actual page content (text chunks, tables, code) relevant to the query, \ + ready for RAG or fact-checking. Supports location-aware queries via optional \ + loc_lat, loc_long, loc_city, loc_state, loc_country, etc. for local/POI results. \ + Use when you need substantive content from the web rather than just links and \ + snippets. Authentication via 'brave_api_key' (same as Web Search)." + .to_string() + } +} + +/// Input parameters for the LLM Context API. Snake_case fields map to Brave's JSON body +/// and optional X-Loc-* headers; validation happens in `validate_params`, clamping in `build_request_body`. +#[derive(Debug, Default, Deserialize)] +struct LlmContextParams { + #[serde(default)] + query: String, + country: Option, + search_lang: Option, + count: Option, + // Context Size Parameters + maximum_number_of_urls: Option, + maximum_number_of_tokens: Option, + maximum_number_of_snippets: Option, + maximum_number_of_tokens_per_url: Option, + maximum_number_of_snippets_per_url: Option, + // Filtering and Local Parameters + context_threshold_mode: Option, + goggles: Option, + // Location-aware query headers + #[serde(rename = "loc_lat")] + loc_lat: Option, + #[serde(rename = "loc_long")] + loc_long: Option, + #[serde(rename = "loc_city")] + loc_city: Option, + #[serde(rename = "loc_state")] + loc_state: Option, + #[serde(rename = "loc_state_name")] + loc_state_name: Option, + #[serde(rename = "loc_country")] + loc_country: Option, + #[serde(rename = "loc_postal_code")] + loc_postal_code: Option, +} + +/// Top-level Brave LLM Context API response: optional grounding (generic/poi/map) and optional sources map. +#[derive(Debug, Deserialize)] +struct BraveLlmContextResponse { + grounding: Option, + sources: Option>, +} + +/// Grounding content by type. See [LLM Context API](https://api-dashboard.search.brave.com/documentation/services/llm-context) and [LLM Context POST](https://api-dashboard.search.brave.com/api-reference/summarizer/llm_context/post). +#[derive(Debug, Deserialize)] +struct Grounding { + /// Main grounding data: array of URL objects with extracted content (text chunks, tables, code). + generic: Option>, + /// Point-of-interest data, sometimes present when local recall is enabled (e.g. via X-Loc-* headers or enable_local). + poi: Option, + /// Map/place results when local recall is enabled. Array of place entries with name, url, title, snippets. + map: Option>, +} + +/// One URL's extracted content in `grounding.generic`: url, title, and text snippets. +#[derive(Clone, Debug, Deserialize)] +struct GenericEntry { + url: Option, + title: Option, + snippets: Option>, +} + +/// Entry shape for `grounding.poi` (single object) and `grounding.map` (array). Present when local recall is active. +#[derive(Debug, Deserialize)] +struct PoiMapEntry { + name: Option, + url: Option, + title: Option, + snippets: Option>, +} + +/// Validate the input parameters against the schema. +fn validate_params(params: &LlmContextParams) -> Result<(), String> { + let trimmed = params.query.trim(); + if trimmed.is_empty() { + return Err("'query' must not be empty or only whitespace".into()); + } + if trimmed.chars().count() > MAX_QUERY_LEN { + return Err(format!( + "'query' exceeds maximum length of {} characters", + MAX_QUERY_LEN + )); + } + let word_count = trimmed.split_whitespace().count(); + if word_count > MAX_QUERY_WORDS { + return Err(format!( + "'query' exceeds maximum of {} words (got {})", + MAX_QUERY_WORDS, word_count + )); + } + + // Validate optional parameters (same style as Web Search tool) + if let Some(ref lang) = params.search_lang { + if !is_valid_lang_code(lang) { + return Err(format!( + "Invalid 'search_lang': expected 2-letter code like 'en', got '{lang}'" + )); + } + } + if let Some(ref country) = params.country { + if !is_valid_country_code(country) { + return Err(format!( + "Invalid 'country': expected 2-letter code like 'US', got '{country}'" + )); + } + } + if let Some(ref mode) = params.context_threshold_mode { + if !is_valid_threshold_mode(mode) { + return Err(format!( + "Invalid 'context_threshold_mode': expected 'strict', 'balanced', 'lenient', or 'disabled', got '{mode}'" + )); + } + } + + if let Some(ref goggles) = params.goggles { + if !is_valid_goggles_value(goggles) { + return Err(format!( + "Invalid 'goggles': expected a non-empty string or a non-empty array of strings (URLs or inline definitions), got '{goggles}'" + )); + } + } + + if let Some(lat) = params.loc_lat { + if !(-90.0..=90.0).contains(&lat) { + return Err(format!( + "Invalid 'loc_lat': must be between -90 and 90 (got {lat})" + )); + } + } + if let Some(long) = params.loc_long { + if !(-180.0..=180.0).contains(&long) { + return Err(format!( + "Invalid 'loc_long': must be between -180 and 180 (got {long})" + )); + } + } + if let Some(ref c) = params.loc_country { + if !is_valid_country_code(c) { + return Err(format!( + "Invalid 'loc_country': expected 2-letter uppercase code like 'US', got '{c}'" + )); + } + } + Ok(()) +} + +/// Entry point: parse, validate, call API, format output. +fn execute_inner(params: &str) -> Result { + let params: LlmContextParams = + serde_json::from_str(params).map_err(|e| format!("Invalid parameters: {e}"))?; + + validate_params(¶ms)?; + preflight_check()?; + + let response_body = call_brave_api(¶ms)?; + let api_response: BraveLlmContextResponse = serde_json::from_str(&response_body) + .map_err(|e| format!("Failed to parse Brave response: {e}"))?; + + format_output(¶ms.query, api_response) +} + +/// Verify the API key is available before making the request. +fn preflight_check() -> Result<(), String> { + if !near::agent::host::secret_exists("brave_api_key") { + return Err("Brave API key not found in secret store. Set it with: \ + ironclaw secret set brave_api_key . \ + Get a key at: https://brave.com/search/api/" + .into()); + } + Ok(()) +} + +/// Call the Brave LLM Context API with retry on transient server errors. +/// +/// Retries on 5xx errors only. 429 (rate limit) is not retried since the WASM +/// sandbox has no sleep primitive and immediate retry would just hit the limit again. +fn call_brave_api(params: &LlmContextParams) -> Result { + let request_body = build_request_body(params)?; + let headers = build_request_headers(params); + + let mut attempt = 0; + let response = loop { + attempt += 1; + + let resp = near::agent::host::http_request( + "POST", + BRAVE_LLM_CONTEXT_ENDPOINT, + &headers.to_string(), + Some(&request_body), + None, + ) + .map_err(|e| format!("HTTP request failed: {e}"))?; + + if resp.status >= 200 && resp.status < 300 { + break resp; + } + + if attempt < MAX_RETRIES && resp.status >= 500 { + near::agent::host::log( + near::agent::host::LogLevel::Warn, + &format!( + "Brave LLM Context API error {} (attempt {}/{}). Retrying...", + resp.status, attempt, MAX_RETRIES + ), + ); + continue; + } + + let error_body = String::from_utf8_lossy(&resp.body); + return Err(format!( + "Brave LLM Context API error (HTTP {}): {}", + resp.status, error_body + )); + }; + + String::from_utf8(response.body).map_err(|e| format!("Invalid UTF-8 response: {e}")) +} + +/// Normalize grounding + sources into a single JSON output. +fn format_output(query: &str, response: BraveLlmContextResponse) -> Result { + let sources = response.sources.unwrap_or_default(); + let grounding = response.grounding; + + let generic = grounding + .as_ref() + .and_then(|g| g.generic.as_deref()) + .unwrap_or_default(); + + let poi = grounding.as_ref().and_then(|g| g.poi.as_ref()); + let map = grounding + .as_ref() + .and_then(|g| g.map.as_deref()) + .unwrap_or_default(); + + // Count snippets from typed data before creating JSON for better performance and type safety. + let generic_snippet_count: usize = generic + .iter() + .map(|e| e.snippets.as_deref().unwrap_or_default().len()) + .sum(); + let poi_snippet_count: usize = poi + .map(|p| p.snippets.as_deref().unwrap_or_default().len()) + .unwrap_or(0); + let map_snippet_count: usize = map + .iter() + .map(|e| e.snippets.as_deref().unwrap_or_default().len()) + .sum(); + let snippet_count = generic_snippet_count + poi_snippet_count + map_snippet_count; + + let entries: Vec = generic + .iter() + .filter_map(|e| { + let url = e.url.as_ref()?; + let title = e.title.as_deref().unwrap_or("Untitled"); + let snippets = e.snippets.as_deref().unwrap_or(&[]); + Some(build_entry_json(url, title, None, snippets, &sources)) + }) + .collect(); + + let poi_output = poi.map(|e| poi_map_entry_to_json(e, &sources)); + + let map_output: Vec = map + .iter() + .map(|e| poi_map_entry_to_json(e, &sources)) + .collect(); + + let mut output = serde_json::json!({ + "query": query, + "url_count": entries.len(), + "snippet_count": snippet_count, + "sources": entries, + }); + + if let Some(poi) = poi_output { + output["poi"] = poi; + } + + if !map_output.is_empty() { + output["map"] = serde_json::json!(map_output); + } + + serde_json::to_string(&output).map_err(|e| format!("Failed to serialize output: {e}")) +} + +/// Build the POST request body as JSON. Clamps numeric fields to API min/max; only includes +/// optional fields when present and valid. +fn build_request_body(params: &LlmContextParams) -> Result, String> { + let count = params + .count + .unwrap_or(DEFAULT_COUNT) + .clamp(MIN_COUNT, MAX_COUNT); + let max_tokens = params + .maximum_number_of_tokens + .unwrap_or(DEFAULT_MAX_TOKENS) + .clamp(MIN_TOKENS, MAX_TOKENS); + let max_urls = params + .maximum_number_of_urls + .unwrap_or(DEFAULT_MAX_URLS) + .clamp(MIN_URLS, MAX_URLS); + let max_snippets = params + .maximum_number_of_snippets + .unwrap_or(DEFAULT_MAX_SNIPPETS) + .clamp(MIN_SNIPPETS, MAX_SNIPPETS); + let max_tokens_per_url = params + .maximum_number_of_tokens_per_url + .unwrap_or(DEFAULT_MAX_TOKENS_PER_URL) + .clamp(MIN_TOKENS_PER_URL, MAX_TOKENS_PER_URL); + let max_snippets_per_url = params + .maximum_number_of_snippets_per_url + .unwrap_or(DEFAULT_SNIPPETS_PER_URL) + .clamp(MIN_SNIPPETS_PER_URL, MAX_SNIPPETS_PER_URL); + + let mut body = serde_json::Map::new(); + body.insert( + "q".to_string(), + serde_json::Value::String(params.query.trim().to_string()), + ); + + // Insert number fields + let number_fields: [(&str, u32); 6] = [ + ("count", count), + ("maximum_number_of_tokens", max_tokens), + ("maximum_number_of_urls", max_urls), + ("maximum_number_of_snippets", max_snippets), + ("maximum_number_of_tokens_per_url", max_tokens_per_url), + ("maximum_number_of_snippets_per_url", max_snippets_per_url), + ]; + for (key, value) in number_fields { + body.insert( + key.to_string(), + serde_json::Value::Number(serde_json::Number::from(value)), + ); + } + + // Optional body fields: + let optional_body_strings: [(&str, Option); 3] = [ + ("country", params.country.clone()), + ("search_lang", params.search_lang.clone()), + ( + "context_threshold_mode", + params.context_threshold_mode.clone(), + ), + ]; + for (key, value) in optional_body_strings { + if let Some(v) = value { + body.insert(key.to_string(), serde_json::Value::String(v)); + } + } + if let Some(goggles) = params.goggles.clone() { + body.insert("goggles".to_string(), goggles); + } + + serde_json::to_vec(&serde_json::Value::Object(body)) + .map_err(|e| format!("Failed to serialize request body: {e}")) +} + +/// Build HTTP request headers: Accept, Content-Type, User-Agent, and optional X-Loc-* +/// for location-aware queries. API key is injected by the host (same as Web Search). +fn build_request_headers(params: &LlmContextParams) -> serde_json::Value { + let mut map = serde_json::Map::new(); + map.insert( + "Accept".to_string(), + serde_json::Value::String("application/json".to_string()), + ); + map.insert( + "Content-Type".to_string(), + serde_json::Value::String("application/json".to_string()), + ); + map.insert( + "User-Agent".to_string(), + serde_json::Value::String("IronClaw-LlmContext-Tool/0.1".to_string()), + ); + + // Location-aware headers: (X-Loc-* name, optional value from params) + let loc_headers: [(&str, Option); 7] = [ + ("X-Loc-Lat", params.loc_lat.map(|v| v.to_string())), + ("X-Loc-Long", params.loc_long.map(|v| v.to_string())), + ("X-Loc-City", params.loc_city.clone()), + ("X-Loc-State", params.loc_state.clone()), + ("X-Loc-State-Name", params.loc_state_name.clone()), + ("X-Loc-Country", params.loc_country.clone()), + ("X-Loc-Postal-Code", params.loc_postal_code.clone()), + ]; + for (header, value) in loc_headers { + if let Some(v) = value { + map.insert(header.to_string(), serde_json::Value::String(v)); + } + } + + serde_json::Value::Object(map) +} + +/// Builds a JSON object for a search result entry. +fn build_entry_json( + url: &str, + title: &str, + name: Option<&str>, + snippets: &[String], + sources: &serde_json::Map, +) -> serde_json::Value { + let hostname = sources + .get(url) + .and_then(|v| v.get("hostname")) + .and_then(|v| v.as_str()) + .map(String::from) + .unwrap_or_else(|| extract_hostname(url).unwrap_or_default()); + + let age_str = sources + .get(url) + .and_then(|v| v.get("age")) + .and_then(|v| v.as_array()) + .and_then(|a| a.first()) + .and_then(|v| v.as_str()); + + let mut entry = serde_json::json!({ + "url": url, + "title": title, + "hostname": hostname, + "snippets": snippets, + }); + + if let Some(name) = name { + entry["name"] = serde_json::json!(name); + } + if let Some(age) = age_str { + entry["age"] = serde_json::json!(age); + } + + entry +} + +/// Build a JSON object for a POI or map entry (name, url, title, hostname, snippets, age when available). +fn poi_map_entry_to_json( + e: &PoiMapEntry, + sources: &serde_json::Map, +) -> serde_json::Value { + let url = e.url.as_deref().unwrap_or_default(); + let title = e.title.as_deref().unwrap_or("Untitled"); + let name = e.name.as_deref(); + let snippets = e.snippets.as_deref().unwrap_or(&[]); + build_entry_json(url, title, name, snippets, sources) +} + +/// Extract hostname from a URL string (no URL parser dependency). Handles http(s) and strips port. +fn extract_hostname(url: &str) -> Option { + let after_scheme = url + .strip_prefix("https://") + .or_else(|| url.strip_prefix("http://"))?; + let host = after_scheme.split('/').next()?; + let host = host.split(':').next()?; + if host.is_empty() { + None + } else { + Some(host.to_string()) + } +} + +/// Validate a 2-letter language code (e.g. "en", "de"). +fn is_valid_lang_code(s: &str) -> bool { + s.len() == 2 && s.bytes().all(|b| b.is_ascii_lowercase()) +} + +/// Validate a 2-letter country code (e.g. "US", "DE"). +fn is_valid_country_code(s: &str) -> bool { + s.len() == 2 && s.bytes().all(|b| b.is_ascii_uppercase()) +} + +/// Validate context_threshold_mode: strict, balanced, lenient, or disabled. +fn is_valid_threshold_mode(s: &str) -> bool { + VALID_THRESHOLD_MODES.contains(&s) +} + +/// Goggles must be a non-empty string or a non-empty array of strings (URLs or inline definitions). +fn is_valid_goggles_value(v: &serde_json::Value) -> bool { + match v { + serde_json::Value::String(s) => !s.is_empty(), + serde_json::Value::Array(a) => { + !a.is_empty() + && a.iter() + .all(|e| matches!(e, serde_json::Value::String(s) if !s.is_empty())) + } + _ => false, + } +} + +// Schema must remain in sync with the MIN_*, DEFAULT_*, and MAX_* constants. +const SCHEMA: &str = r#"{ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query; returns pre-extracted web content (text, tables, code) for grounding LLM answers", + "minLength": 1, + "maxLength": 400 + }, + "count": { + "type": "integer", + "description": "Maximum number of search results to consider (1-50, default 20)", + "minimum": 1, + "maximum": 50, + "default": 20 + }, + "country": { + "type": "string", + "description": "2-letter uppercase country code (e.g. 'US', 'DE')" + }, + "search_lang": { + "type": "string", + "description": "2-letter lowercase language code for results (e.g. 'en', 'de')" + }, + "maximum_number_of_tokens": { + "type": "integer", + "description": "Approximate max tokens in returned context (1024-32768, default 8192)", + "minimum": 1024, + "maximum": 32768, + "default": 8192 + }, + "maximum_number_of_urls": { + "type": "integer", + "description": "Maximum URLs to include (1-50, default 20)", + "minimum": 1, + "maximum": 50, + "default": 20 + }, + "maximum_number_of_snippets": { + "type": "integer", + "description": "Maximum snippets across all URLs (1-100, default 50)", + "minimum": 1, + "maximum": 100, + "default": 50 + }, + "maximum_number_of_tokens_per_url": { + "type": "integer", + "description": "Max tokens per URL (512-8192, default 4096)", + "minimum": 512, + "maximum": 8192, + "default": 4096 + }, + "maximum_number_of_snippets_per_url": { + "type": "integer", + "description": "Max snippets per URL (1-100, default 50)", + "minimum": 1, + "maximum": 100, + "default": 50 + }, + "context_threshold_mode": { + "type": "string", + "description": "Relevance filter: 'strict' (fewer, more relevant), 'balanced', 'lenient', or 'disabled'", + "enum": ["strict", "balanced", "lenient", "disabled"] + }, + "loc_lat": { + "type": "number", + "description": "Latitude for location-aware queries (-90 to 90). Use with loc_long or place-name headers for local/POI results." + }, + "loc_long": { + "type": "number", + "description": "Longitude for location-aware queries (-180 to 180). Use with loc_lat or place-name headers for local/POI results." + }, + "loc_city": { + "type": "string", + "description": "City name for location-aware queries (e.g. 'San Francisco')" + }, + "loc_state": { + "type": "string", + "description": "State/region code for location-aware queries (e.g. 'CA', ISO 3166-2)" + }, + "loc_state_name": { + "type": "string", + "description": "State/region full name for location-aware queries" + }, + "loc_country": { + "type": "string", + "description": "2-letter uppercase country code for location headers (e.g. 'US'). Enables local recall for queries like 'coffee shops near me'." + }, + "loc_postal_code": { + "type": "string", + "description": "Postal code for location-aware queries" + }, + "goggles": { + "description": "Custom ranking/filtering: URL to a Goggle file, inline Goggles rules, or array of URLs/inline strings. Restrict or boost sources (e.g. trusted domains). See https://api-dashboard.search.brave.com/documentation/resources/goggles", + "oneOf": [ + { "type": "string", "minLength": 1 }, + { "type": "array", "items": { "type": "string", "minLength": 1 }, "minItems": 1 } + ] + } + }, + "required": ["query"], + "additionalProperties": false +}"#; + +export!(LlmContextTool); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_extract_hostname() { + assert_eq!( + extract_hostname("https://example.com/path"), + Some("example.com".into()) + ); + assert_eq!( + extract_hostname("http://example.com"), + Some("example.com".into()) + ); + assert_eq!( + extract_hostname("http://host:8080/path"), + Some("host".into()) + ); + assert_eq!( + extract_hostname("https://sub.example.com:443/"), + Some("sub.example.com".into()) + ); + assert_eq!(extract_hostname("https://"), None); + assert_eq!(extract_hostname("https:///path"), None); + assert_eq!(extract_hostname("ftp://example.com"), None); + assert_eq!(extract_hostname("example.com"), None); + assert_eq!(extract_hostname(""), None); + } + + #[test] + fn test_is_valid_lang_code() { + assert!(is_valid_lang_code("en")); + assert!(!is_valid_lang_code("EN")); + assert!(!is_valid_lang_code("eng")); + } + + #[test] + fn test_is_valid_country_code() { + assert!(is_valid_country_code("US")); + assert!(!is_valid_country_code("us")); + assert!(!is_valid_country_code("USA")); + } + + #[test] + fn test_is_valid_threshold_mode() { + assert!(is_valid_threshold_mode("strict")); + assert!(is_valid_threshold_mode("balanced")); + assert!(is_valid_threshold_mode("lenient")); + assert!(is_valid_threshold_mode("disabled")); + assert!(!is_valid_threshold_mode("invalid")); + } + + fn params_minimal() -> LlmContextParams { + LlmContextParams { + query: "rust async".to_string(), + ..Default::default() + } + } + + #[test] + fn test_validate_params_accepts_minimal() { + let params = params_minimal(); + assert!(validate_params(¶ms).is_ok()); + } + + #[test] + fn test_validate_params_rejects_invalid() { + // Empty query + let mut p = params_minimal(); + p.query = "".to_string(); + assert!(validate_params(&p).is_err()); + + // Query too long + p.query = "a".repeat(MAX_QUERY_LEN + 1); + assert!(validate_params(&p).is_err()); + + // Too many words + p.query = (0..MAX_QUERY_WORDS + 1) + .map(|i| format!("w{i}")) + .collect::>() + .join(" "); + assert!(validate_params(&p).is_err()); + + // Invalid search_lang (must be 2-letter lowercase) + p = params_minimal(); + p.search_lang = Some("EN".to_string()); + assert!(validate_params(&p).is_err()); + + // Invalid country (must be 2-letter uppercase) + p = params_minimal(); + p.country = Some("us".to_string()); + assert!(validate_params(&p).is_err()); + + // Invalid context_threshold_mode + p = params_minimal(); + p.context_threshold_mode = Some("invalid".to_string()); + assert!(validate_params(&p).is_err()); + + // Invalid loc_lat (out of range) + p = params_minimal(); + p.loc_lat = Some(91.0); + assert!(validate_params(&p).is_err()); + + // Invalid loc_long (out of range) + p = params_minimal(); + p.loc_long = Some(-181.0); + assert!(validate_params(&p).is_err()); + + // Invalid loc_country + p = params_minimal(); + p.loc_country = Some("usa".to_string()); + assert!(validate_params(&p).is_err()); + + // Invalid goggles (empty string) + p = params_minimal(); + p.goggles = Some(serde_json::Value::String(String::new())); + assert!(validate_params(&p).is_err()); + } + + #[test] + fn test_build_request_body_minimal() { + let params = params_minimal(); + let body = build_request_body(¶ms).unwrap(); + let obj: serde_json::Map = + serde_json::from_slice(&body).unwrap(); + assert_eq!(obj.get("q").and_then(|v| v.as_str()), Some("rust async")); + assert_eq!(obj.get("count").and_then(|v| v.as_u64()), Some(20)); + assert_eq!( + obj.get("maximum_number_of_tokens").and_then(|v| v.as_u64()), + Some(8192) + ); + assert!(!obj.contains_key("country")); + assert!(!obj.contains_key("context_threshold_mode")); + } + + #[test] + fn test_build_request_body_full() { + let params = LlmContextParams { + query: "python asyncio".to_string(), + count: Some(10), + country: Some("US".to_string()), + search_lang: Some("en".to_string()), + maximum_number_of_tokens: Some(4096), + maximum_number_of_urls: Some(10), + maximum_number_of_snippets: Some(25), + maximum_number_of_tokens_per_url: Some(2048), + maximum_number_of_snippets_per_url: Some(25), + context_threshold_mode: Some("strict".to_string()), + ..Default::default() + }; + let body = build_request_body(¶ms).unwrap(); + let obj: serde_json::Map = + serde_json::from_slice(&body).unwrap(); + assert_eq!( + obj.get("q").and_then(|v| v.as_str()), + Some("python asyncio") + ); + assert_eq!(obj.get("count").and_then(|v| v.as_u64()), Some(10)); + assert_eq!(obj.get("country").and_then(|v| v.as_str()), Some("US")); + assert_eq!(obj.get("search_lang").and_then(|v| v.as_str()), Some("en")); + assert_eq!( + obj.get("maximum_number_of_tokens").and_then(|v| v.as_u64()), + Some(4096) + ); + assert_eq!( + obj.get("context_threshold_mode").and_then(|v| v.as_str()), + Some("strict") + ); + } + + #[test] + fn test_build_request_headers_with_location() { + let params = LlmContextParams { + query: "coffee shops".to_string(), + loc_lat: Some(37.7749), + loc_long: Some(-122.4194), + loc_city: Some("San Francisco".to_string()), + loc_state: Some("CA".to_string()), + loc_state_name: Some("California".to_string()), + loc_country: Some("US".to_string()), + loc_postal_code: Some("94102".to_string()), + ..Default::default() + }; + let headers = build_request_headers(¶ms); + let obj = headers.as_object().unwrap(); + assert_eq!( + obj.get("Accept").and_then(|v| v.as_str()), + Some("application/json") + ); + assert_eq!( + obj.get("X-Loc-Lat").and_then(|v| v.as_str()), + Some("37.7749") + ); + assert_eq!( + obj.get("X-Loc-Long").and_then(|v| v.as_str()), + Some("-122.4194") + ); + assert_eq!( + obj.get("X-Loc-City").and_then(|v| v.as_str()), + Some("San Francisco") + ); + assert_eq!(obj.get("X-Loc-State").and_then(|v| v.as_str()), Some("CA")); + assert_eq!( + obj.get("X-Loc-State-Name").and_then(|v| v.as_str()), + Some("California") + ); + assert_eq!( + obj.get("X-Loc-Country").and_then(|v| v.as_str()), + Some("US") + ); + assert_eq!( + obj.get("X-Loc-Postal-Code").and_then(|v| v.as_str()), + Some("94102") + ); + } + + #[test] + fn test_build_request_headers_no_location() { + let params = params_minimal(); + let headers = build_request_headers(¶ms); + let obj = headers.as_object().unwrap(); + assert_eq!( + obj.get("Accept").and_then(|v| v.as_str()), + Some("application/json") + ); + assert_eq!( + obj.get("Content-Type").and_then(|v| v.as_str()), + Some("application/json") + ); + assert!(obj.get("User-Agent").is_some()); + assert!(obj.get("X-Loc-Lat").is_none()); + assert!(obj.get("X-Loc-Country").is_none()); + } + + #[test] + fn test_build_request_body_with_goggles_string() { + let mut params = params_minimal(); + params.query = "rust programming".to_string(); + params.goggles = Some(serde_json::Value::String( + "https://raw.githubusercontent.com/brave/goggles-quickstart/main/goggles/tech_blogs.goggle" + .to_string(), + )); + let body = build_request_body(¶ms).unwrap(); + let obj: serde_json::Map = + serde_json::from_slice(&body).unwrap(); + assert_eq!( + obj.get("goggles").and_then(|v| v.as_str()), + Some("https://raw.githubusercontent.com/brave/goggles-quickstart/main/goggles/tech_blogs.goggle") + ); + } + + #[test] + fn test_build_request_body_with_goggles_array() { + let mut params = params_minimal(); + params.query = "web development".to_string(); + params.goggles = Some(serde_json::json!([ + "https://example.com/goggle1.goggle", + "$boost=3,site=dev.to" + ])); + let body = build_request_body(¶ms).unwrap(); + let obj: serde_json::Map = + serde_json::from_slice(&body).unwrap(); + let arr = obj.get("goggles").and_then(|v| v.as_array()).unwrap(); + assert_eq!(arr.len(), 2); + assert_eq!(arr[0].as_str(), Some("https://example.com/goggle1.goggle")); + assert_eq!(arr[1].as_str(), Some("$boost=3,site=dev.to")); + } + + #[test] + fn test_is_valid_goggles_value() { + assert!(is_valid_goggles_value(&serde_json::Value::String( + "https://x.com/a.goggle".to_string() + ))); + assert!(is_valid_goggles_value(&serde_json::json!([ + "https://a.com", + "$boost,site=dev.to" + ]))); + assert!(!is_valid_goggles_value(&serde_json::Value::String( + "".to_string() + ))); + assert!(!is_valid_goggles_value(&serde_json::Value::Array(vec![]))); + assert!(!is_valid_goggles_value(&serde_json::Value::Bool(true))); + } + + #[test] + fn test_parse_response() { + let body = r#"{ + "grounding": { + "generic": [ + { + "url": "https://example.com/page", + "title": "Example Page", + "snippets": ["First snippet.", "Second snippet."] + } + ] + }, + "sources": { + "https://example.com/page": { + "title": "Example Page", + "hostname": "example.com", + "age": ["2024-01-15", "380 days ago"] + } + } + }"#; + let r: BraveLlmContextResponse = serde_json::from_str(body).unwrap(); + let generic = r.grounding.unwrap().generic.unwrap(); + assert_eq!(generic.len(), 1); + assert_eq!(generic[0].url.as_deref(), Some("https://example.com/page")); + assert_eq!(generic[0].title.as_deref(), Some("Example Page")); + assert_eq!(generic[0].snippets.as_ref().unwrap().len(), 2); + let sources = r.sources.unwrap(); + let meta = sources.get("https://example.com/page").unwrap(); + assert_eq!( + meta.get("hostname").and_then(|v| v.as_str()), + Some("example.com") + ); + } + + #[test] + fn test_parse_response_with_poi_and_map() { + let body = r#"{ + "grounding": { + "generic": [{"url": "https://example.com/page", "title": "Example", "snippets": []}], + "poi": { + "name": "Business Name", + "url": "https://business.com", + "title": "Title of business.com website", + "snippets": ["Business details."] + }, + "map": [ + { + "name": "Place Name", + "url": "https://place.com", + "title": "Title of place.com", + "snippets": ["Place information."] + } + ] + }, + "sources": { + "https://business.com": {"title": "Business Name", "hostname": "business.com", "age": null}, + "https://place.com": {"title": "Place", "hostname": "place.com", "age": null} + } + }"#; + let r: BraveLlmContextResponse = serde_json::from_str(body).unwrap(); + let g = r.grounding.as_ref().unwrap(); + assert_eq!(g.generic.as_ref().unwrap().len(), 1); + let poi = g.poi.as_ref().unwrap(); + assert_eq!(poi.name.as_deref(), Some("Business Name")); + assert_eq!(poi.url.as_deref(), Some("https://business.com")); + assert_eq!(poi.snippets.as_ref().unwrap().len(), 1); + let map = g.map.as_ref().unwrap(); + assert_eq!(map.len(), 1); + assert_eq!(map[0].name.as_deref(), Some("Place Name")); + assert_eq!(map[0].url.as_deref(), Some("https://place.com")); + } + + #[test] + fn test_poi_map_entry_to_json() { + let e = PoiMapEntry { + name: Some("Cafe Example".to_string()), + url: Some("https://cafe.example.com".to_string()), + title: Some("Cafe Example - Coffee".to_string()), + snippets: Some(vec!["Best coffee in town.".to_string()]), + }; + let mut sources = serde_json::Map::new(); + sources.insert( + "https://cafe.example.com".to_string(), + serde_json::json!({"hostname": "cafe.example.com", "age": ["2024-06-01"]}), + ); + let out = poi_map_entry_to_json(&e, &sources); + assert_eq!( + out.get("name").and_then(|v| v.as_str()), + Some("Cafe Example") + ); + assert_eq!( + out.get("url").and_then(|v| v.as_str()), + Some("https://cafe.example.com") + ); + assert_eq!( + out.get("hostname").and_then(|v| v.as_str()), + Some("cafe.example.com") + ); + assert_eq!(out.get("age").and_then(|v| v.as_str()), Some("2024-06-01")); + let snippets = out.get("snippets").and_then(|s| s.as_array()).unwrap(); + assert_eq!(snippets.len(), 1); + assert_eq!(snippets[0].as_str(), Some("Best coffee in town.")); + } + + #[test] + fn test_build_request_body_clamps_below_min() { + let mut params = params_minimal(); + params.count = Some(0); + params.maximum_number_of_tokens = Some(100); + params.maximum_number_of_urls = Some(0); + params.maximum_number_of_snippets = Some(0); + params.maximum_number_of_tokens_per_url = Some(1); + params.maximum_number_of_snippets_per_url = Some(0); + + let body = build_request_body(¶ms).unwrap(); + let obj: serde_json::Map = + serde_json::from_slice(&body).unwrap(); + + assert_eq!(obj["count"].as_u64(), Some(MIN_COUNT as u64)); + assert_eq!( + obj["maximum_number_of_tokens"].as_u64(), + Some(MIN_TOKENS as u64) + ); + assert_eq!( + obj["maximum_number_of_urls"].as_u64(), + Some(MIN_URLS as u64) + ); + assert_eq!( + obj["maximum_number_of_snippets"].as_u64(), + Some(MIN_SNIPPETS as u64) + ); + assert_eq!( + obj["maximum_number_of_tokens_per_url"].as_u64(), + Some(MIN_TOKENS_PER_URL as u64) + ); + assert_eq!( + obj["maximum_number_of_snippets_per_url"].as_u64(), + Some(MIN_SNIPPETS_PER_URL as u64) + ); + } + + #[test] + fn test_build_request_body_clamps_above_max() { + let mut params = params_minimal(); + params.count = Some(999); + params.maximum_number_of_tokens = Some(999_999); + params.maximum_number_of_urls = Some(999); + params.maximum_number_of_snippets = Some(999); + params.maximum_number_of_tokens_per_url = Some(999_999); + params.maximum_number_of_snippets_per_url = Some(999); + + let body = build_request_body(¶ms).unwrap(); + let obj: serde_json::Map = + serde_json::from_slice(&body).unwrap(); + + assert_eq!(obj["count"].as_u64(), Some(MAX_COUNT as u64)); + assert_eq!( + obj["maximum_number_of_tokens"].as_u64(), + Some(MAX_TOKENS as u64) + ); + assert_eq!( + obj["maximum_number_of_urls"].as_u64(), + Some(MAX_URLS as u64) + ); + assert_eq!( + obj["maximum_number_of_snippets"].as_u64(), + Some(MAX_SNIPPETS as u64) + ); + assert_eq!( + obj["maximum_number_of_tokens_per_url"].as_u64(), + Some(MAX_TOKENS_PER_URL as u64) + ); + assert_eq!( + obj["maximum_number_of_snippets_per_url"].as_u64(), + Some(MAX_SNIPPETS_PER_URL as u64) + ); + } + + #[test] + fn test_build_entry_json_missing_source() { + let sources = serde_json::Map::new(); + let entry = build_entry_json( + "https://unknown.com/page", + "Title", + None, + &["snippet".to_string()], + &sources, + ); + assert_eq!( + entry.get("hostname").and_then(|v| v.as_str()), + Some("unknown.com") + ); + assert!(entry.get("age").is_none()); + } + + #[test] + fn test_build_entry_json_with_name() { + let sources = serde_json::Map::new(); + let entry = build_entry_json( + "https://example.com", + "Title", + Some("My Place"), + &[], + &sources, + ); + assert_eq!(entry.get("name").and_then(|v| v.as_str()), Some("My Place")); + } + + #[test] + fn test_parse_empty_grounding_response() { + let body = r#"{"grounding": null, "sources": null}"#; + let r: BraveLlmContextResponse = serde_json::from_str(body).unwrap(); + assert!(r.grounding.is_none()); + assert!(r.sources.is_none()); + } + + #[test] + fn test_parse_empty_generic_array() { + let body = r#"{"grounding": {"generic": []}, "sources": {}}"#; + let r: BraveLlmContextResponse = serde_json::from_str(body).unwrap(); + assert!(r.grounding.unwrap().generic.unwrap().is_empty()); + } + + #[test] + fn test_format_output_empty_response() { + let response = BraveLlmContextResponse { + grounding: None, + sources: None, + }; + let result = format_output("test query", response).unwrap(); + let parsed: serde_json::Value = serde_json::from_str(&result).unwrap(); + assert_eq!(parsed["query"].as_str(), Some("test query")); + assert_eq!(parsed["url_count"].as_u64(), Some(0)); + assert_eq!(parsed["snippet_count"].as_u64(), Some(0)); + assert!(parsed["sources"].as_array().unwrap().is_empty()); + assert!(parsed.get("poi").is_none()); + assert!(parsed.get("map").is_none()); + } + + #[test] + fn test_format_output_with_generic_entries() { + let response = BraveLlmContextResponse { + grounding: Some(Grounding { + generic: Some(vec![ + GenericEntry { + url: Some("https://example.com".to_string()), + title: Some("Example".to_string()), + snippets: Some(vec!["s1".to_string(), "s2".to_string()]), + }, + GenericEntry { + url: None, + title: Some("No URL".to_string()), + snippets: None, + }, + ]), + poi: None, + map: None, + }), + sources: Some({ + let mut m = serde_json::Map::new(); + m.insert( + "https://example.com".to_string(), + serde_json::json!({"hostname": "example.com", "age": ["2024-01-01"]}), + ); + m + }), + }; + let result = format_output("test", response).unwrap(); + let parsed: serde_json::Value = serde_json::from_str(&result).unwrap(); + assert_eq!(parsed["url_count"].as_u64(), Some(1)); + assert_eq!(parsed["snippet_count"].as_u64(), Some(2)); + let first = &parsed["sources"][0]; + assert_eq!(first["hostname"].as_str(), Some("example.com")); + assert_eq!(first["age"].as_str(), Some("2024-01-01")); + } + + #[test] + fn test_format_output_with_poi_and_map() { + let response = BraveLlmContextResponse { + grounding: Some(Grounding { + generic: Some(vec![]), + poi: Some(PoiMapEntry { + name: Some("Coffee Shop".to_string()), + url: Some("https://coffee.com".to_string()), + title: Some("Coffee".to_string()), + snippets: Some(vec!["Great beans.".to_string()]), + }), + map: Some(vec![PoiMapEntry { + name: Some("Place".to_string()), + url: Some("https://place.com".to_string()), + title: Some("Place".to_string()), + snippets: Some(vec!["Info.".to_string(), "More info.".to_string()]), + }]), + }), + sources: None, + }; + let result = format_output("coffee", response).unwrap(); + let parsed: serde_json::Value = serde_json::from_str(&result).unwrap(); + assert_eq!(parsed["snippet_count"].as_u64(), Some(3)); + assert_eq!(parsed["poi"]["name"].as_str(), Some("Coffee Shop")); + assert_eq!(parsed["map"].as_array().unwrap().len(), 1); + } + + #[test] + fn test_schema_is_valid_json_and_matches_constants() { + let schema: serde_json::Value = + serde_json::from_str(SCHEMA).expect("SCHEMA must be valid JSON"); + let props = schema["properties"].as_object().unwrap(); + + let count = &props["count"]; + assert_eq!(count["minimum"].as_u64(), Some(MIN_COUNT as u64)); + assert_eq!(count["maximum"].as_u64(), Some(MAX_COUNT as u64)); + assert_eq!(count["default"].as_u64(), Some(DEFAULT_COUNT as u64)); + + let max_tokens = &props["maximum_number_of_tokens"]; + assert_eq!(max_tokens["minimum"].as_u64(), Some(MIN_TOKENS as u64)); + assert_eq!(max_tokens["maximum"].as_u64(), Some(MAX_TOKENS as u64)); + assert_eq!( + max_tokens["default"].as_u64(), + Some(DEFAULT_MAX_TOKENS as u64) + ); + + let max_urls = &props["maximum_number_of_urls"]; + assert_eq!(max_urls["minimum"].as_u64(), Some(MIN_URLS as u64)); + assert_eq!(max_urls["maximum"].as_u64(), Some(MAX_URLS as u64)); + assert_eq!(max_urls["default"].as_u64(), Some(DEFAULT_MAX_URLS as u64)); + + let max_snippets = &props["maximum_number_of_snippets"]; + assert_eq!(max_snippets["minimum"].as_u64(), Some(MIN_SNIPPETS as u64)); + assert_eq!(max_snippets["maximum"].as_u64(), Some(MAX_SNIPPETS as u64)); + assert_eq!( + max_snippets["default"].as_u64(), + Some(DEFAULT_MAX_SNIPPETS as u64) + ); + + let max_tpu = &props["maximum_number_of_tokens_per_url"]; + assert_eq!(max_tpu["minimum"].as_u64(), Some(MIN_TOKENS_PER_URL as u64)); + assert_eq!(max_tpu["maximum"].as_u64(), Some(MAX_TOKENS_PER_URL as u64)); + assert_eq!( + max_tpu["default"].as_u64(), + Some(DEFAULT_MAX_TOKENS_PER_URL as u64) + ); + + let max_spu = &props["maximum_number_of_snippets_per_url"]; + assert_eq!( + max_spu["minimum"].as_u64(), + Some(MIN_SNIPPETS_PER_URL as u64) + ); + assert_eq!( + max_spu["maximum"].as_u64(), + Some(MAX_SNIPPETS_PER_URL as u64) + ); + assert_eq!( + max_spu["default"].as_u64(), + Some(DEFAULT_SNIPPETS_PER_URL as u64) + ); + + let query = &props["query"]; + assert_eq!(query["maxLength"].as_u64(), Some(MAX_QUERY_LEN as u64)); + } + + #[test] + fn test_validate_params_trimmed_query_within_limit() { + let mut p = params_minimal(); + p.query = format!(" {} ", "a".repeat(MAX_QUERY_LEN - 4)); + assert!( + validate_params(&p).is_ok(), + "trimmed query within limit should pass" + ); + } + + #[test] + fn test_validate_params_trimmed_query_over_limit() { + let mut p = params_minimal(); + p.query = format!(" {} ", "a".repeat(MAX_QUERY_LEN + 1)); + assert!( + validate_params(&p).is_err(), + "trimmed query over limit should fail" + ); + } +} diff --git a/tools-src/web-search/web-search-tool.capabilities.json b/tools-src/web-search/web-search-tool.capabilities.json index bc660aaf..9c2559ab 100644 --- a/tools-src/web-search/web-search-tool.capabilities.json +++ b/tools-src/web-search/web-search-tool.capabilities.json @@ -1,6 +1,41 @@ { "version": "0.2.0", "wit_version": "0.3.0", + "description": "Search the web using Brave Search. Returns titles, URLs, descriptions, and publication dates for matching web pages. Supports filtering by country, language, and freshness. Authentication is handled via the 'brave_api_key' secret injected by the host.", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query to look up on the web" + }, + "count": { + "type": "integer", + "description": "Number of results to return (1-20, default 5)", + "minimum": 1, + "maximum": 20, + "default": 5 + }, + "country": { + "type": "string", + "description": "2-letter uppercase country code to bias results (e.g. 'US', 'DE', 'JP')" + }, + "search_lang": { + "type": "string", + "description": "2-letter lowercase language code for search results (e.g. 'en', 'de', 'fr')" + }, + "ui_lang": { + "type": "string", + "description": "Locale in language-region format (e.g. 'en-US', 'de-DE')" + }, + "freshness": { + "type": "string", + "description": "Filter by discovery time: 'pd' (past day), 'pw' (past week), 'pm' (past month), 'py' (past year), or date range 'YYYY-MM-DDtoYYYY-MM-DD'" + } + }, + "required": ["query"], + "additionalProperties": false + }, "capabilities": { "http": { "allowlist": [