diff --git a/.github/dependabot.yml b/.github/dependabot.yml index bdd4faf..168b680 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -16,11 +16,15 @@ updates: - package-ecosystem: "github-actions" directory: "/" schedule: - interval: "daily" + interval: "weekly" labels: # Mark PRs as CI related change. - T-CI open-pull-requests-limit: 3 + groups: + codeql-action: + patterns: + - "github/codeql-action*" commit-message: prefix: "chore" - include: "scope" \ No newline at end of file + include: "scope" diff --git a/.github/workflows/auto-label-pr.yml b/.github/workflows/auto-label-pr.yml index f95dde7..d91ce04 100644 --- a/.github/workflows/auto-label-pr.yml +++ b/.github/workflows/auto-label-pr.yml @@ -3,6 +3,9 @@ on: # Runs workflow when activity on a PR in the workflow's repository occurs. pull_request_target: +permissions: + contents: read + jobs: auto-label: permissions: @@ -20,7 +23,7 @@ jobs: PR_URL: ${{ github.event.pull_request.html_url }} steps: - - uses: actions/labeler@v6 + - uses: actions/labeler@bf12e9b00b37c5c0ca2b87b79b2daf7891dbda13 # v7 with: # Auto-include paths starting with dot (e.g. .github) dot: true diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bbe6c26..7ffafe1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -12,34 +12,35 @@ env: CARGO_TERM_COLOR: always ARTIFACT_DIR: release-artifacts +permissions: + contents: read + jobs: commit-lint: name: Lint Commit Messages runs-on: ubuntu-latest if: github.event_name == 'pull_request' steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 - name: Setup Node.js - uses: actions/setup-node@v6 + uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7 with: node-version: '22' - name: Install commitlint - run: | - npm install --save-dev @commitlint/cli @commitlint/config-conventional - echo "module.exports = {extends: ['@commitlint/config-conventional']}" > commitlint.config.js - + run: npm install --global @commitlint/cli@20.4.3 @commitlint/config-conventional@20.4.3 + - name: Lint commit messages - run: npx commitlint --from ${{ github.event.pull_request.base.sha }} --to ${{ github.event.pull_request.head.sha }} --verbose + run: commitlint --extends @commitlint/config-conventional --from ${{ github.event.pull_request.base.sha }} --to ${{ github.event.pull_request.head.sha }} --verbose fmt: name: Code Formatting runs-on: ubuntu-latest steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install Rust fmt run: rustup toolchain install nightly --component rustfmt @@ -51,32 +52,40 @@ jobs: name: Lint with Clippy runs-on: ubuntu-latest steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install Rust - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - - name: Run clippy + - name: Run clippy (all features) run: cargo clippy --all-targets --all-features -- -D warnings + - name: Run clippy (all features except local) + run: | + FEATURES=$(cargo metadata --no-deps --format-version 1 \ + | jq -r '[.packages[] | select(.name == "rmcp") | .features | keys[] + | select(startswith("__") | not) + | select(. != "local")] | join(",")') + cargo clippy --package rmcp --all-targets --no-default-features --features "$FEATURES" -- -D warnings + semver: name: SemVer Check runs-on: ubuntu-latest if: github.event_name == 'pull_request' steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 - name: Install Rust - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - name: Install cargo-semver-checks - uses: taiki-e/install-action@v2 + uses: taiki-e/install-action@3f74d7c16a4242f1c95561e98edc25d36adb4375 # v2.87.12 with: tool: cargo-semver-checks @@ -117,7 +126,7 @@ jobs: runs-on: ubuntu-latest if: github.event_name == 'pull_request' steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 @@ -125,12 +134,12 @@ jobs: # to be installed (it does not need to be the default; the tool invokes it # via `cargo +nightly`). - name: Install Rust - uses: dtolnay/rust-toolchain@nightly + uses: dtolnay/rust-toolchain@7c8d7d138f5c09cef361f8214cf96882cd029cdb # nightly - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - name: Install cargo-public-api - uses: taiki-e/install-action@v2 + uses: taiki-e/install-action@3f74d7c16a4242f1c95561e98edc25d36adb4375 # v2.87.12 with: tool: cargo-public-api @@ -170,10 +179,24 @@ jobs: - name: Check rmcp (all features except local) run: | - FEATURES=$(cargo metadata --no-deps --format-version 1 \ - | jq -r '[.packages[] | select(.name == "rmcp") | .features | keys[] - | select(startswith("__") | not) - | select(. != "local")] | join(",")') + BASE_SHA=${{ github.event.pull_request.base.sha }} + # `cargo public-api diff` builds both revisions with the same feature + # set, so a feature that exists on only one side (e.g. a feature added + # or removed by this PR) would fail to build the other revision. Diff + # only the features present in BOTH the base and the head; features + # unique to one side are necessarily pure additions/removals, which the + # release-type deny flags already govern. + list_features() { + cargo metadata --no-deps --format-version 1 --manifest-path "$1/Cargo.toml" \ + | jq -r '.packages[] | select(.name == "rmcp") | .features | keys[] + | select(startswith("__") | not) + | select(. != "local")' + } + list_features "." | sort -u > "$RUNNER_TEMP/head_features" + git worktree add --detach "$RUNNER_TEMP/rmcp-base" "$BASE_SHA" + list_features "$RUNNER_TEMP/rmcp-base" | sort -u > "$RUNNER_TEMP/base_features" + git worktree remove --force "$RUNNER_TEMP/rmcp-base" + FEATURES=$(comm -12 "$RUNNER_TEMP/head_features" "$RUNNER_TEMP/base_features" | paste -sd, -) cargo public-api \ --package rmcp \ --features "$FEATURES" \ @@ -181,33 +204,57 @@ jobs: diff \ $DENY \ --force \ - ${{ github.event.pull_request.base.sha }}..${{ github.sha }} + "$BASE_SHA"..${{ github.sha }} spelling: name: spell check with typos runs-on: ubuntu-latest steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Spell Check Repo - uses: crate-ci/typos@master + uses: crate-ci/typos@d43b6c087ac471e2ea7b8af622ff15f05c0c365b # master + + msrv: + name: Check MSRV + runs-on: ubuntu-latest + permissions: + contents: read + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Read MSRV from Cargo.toml + run: | + MSRV=$(cargo metadata --no-deps --format-version 1 \ + | jq -r '.packages[] | select(.name == "rmcp") | .rust_version') + echo "MSRV=$MSRV" >> "$GITHUB_ENV" + + - name: Install MSRV Rust + uses: dtolnay/rust-toolchain@6c977a6ca4077a0ceb28ffbe03f59d46e9ac8772 # master + with: + toolchain: ${{ env.MSRV }} + + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + + - name: Check default workspace members with MSRV + run: cargo +${{ env.MSRV }} check --all-targets --all-features test: name: Run Tests runs-on: ubuntu-latest steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 # install nodejs - name: Setup Node.js - uses: actions/setup-node@v6 + uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7 with: node-version: '22' - name: Install uv - uses: astral-sh/setup-uv@v7 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 - name: Install Rust - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - name: Set up Python run: uv python install @@ -215,7 +262,7 @@ jobs: - name: Create venv for python run: uv venv - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - name: Run tests run: cargo test --all-features @@ -224,19 +271,19 @@ jobs: name: Run Tests (no local feature) runs-on: ubuntu-latest steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 # install nodejs - name: Setup Node.js - uses: actions/setup-node@v6 + uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7 with: node-version: '22' - name: Install uv - uses: astral-sh/setup-uv@v7 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 - name: Install Rust - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - name: Set up Python run: uv python install @@ -244,7 +291,7 @@ jobs: - name: Create venv for python run: uv venv - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - name: Run tests without local feature run: | @@ -258,21 +305,21 @@ jobs: name: Code Coverage runs-on: ubuntu-latest permissions: - contents: write + contents: read steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 # install nodejs - name: Setup Node.js - uses: actions/setup-node@v6 + uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7 with: node-version: '22' - name: Install uv - uses: astral-sh/setup-uv@v7 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 - name: Install Rust - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - name: Set up Python run: uv python install @@ -280,7 +327,7 @@ jobs: - name: Create venv for python run: uv venv - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - name: Install cargo-llvm-cov run: cargo install cargo-llvm-cov @@ -295,19 +342,19 @@ jobs: name: Example test runs-on: ubuntu-latest steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 # install nodejs - name: Setup Node.js - uses: actions/setup-node@v6 + uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7 with: node-version: '22' - name: Install uv - uses: astral-sh/setup-uv@v7 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 - name: Install Rust - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - name: Set up Python run: uv python install @@ -315,7 +362,7 @@ jobs: - name: Create venv for python run: uv venv - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - name: Add target WASI preview 2 run: | @@ -351,12 +398,12 @@ jobs: name: Security Audit runs-on: ubuntu-latest steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install Rust - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - name: Install cargo-audit run: cargo install cargo-audit @@ -368,12 +415,12 @@ jobs: name: Generate Documentation runs-on: ubuntu-latest steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install Rust - uses: dtolnay/rust-toolchain@nightly + uses: dtolnay/rust-toolchain@7c8d7d138f5c09cef361f8214cf96882cd029cdb # nightly - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - name: Generate documentation run: | @@ -394,7 +441,7 @@ jobs: # This happened recently in the attack on `tj-actions/changed-files`, but # has happened many times before as well. - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Update Rust run: | diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 0ae4e63..aa77b0a 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -8,6 +8,9 @@ on: schedule: - cron: '0 0 * * 1' # Weekly on Monday +permissions: + contents: read + jobs: analyze: name: Analyze @@ -21,16 +24,16 @@ jobs: language: [rust, javascript-typescript, python, actions] steps: - name: Checkout repository - uses: actions/checkout@v7 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Initialize CodeQL - uses: github/codeql-action/init@v4 + uses: github/codeql-action/init@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 with: languages: ${{ matrix.language }} config-file: ./.github/codeql/codeql-config.yml - name: Autobuild - uses: github/codeql-action/autobuild@v4 + uses: github/codeql-action/autobuild@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@v4 + uses: github/codeql-action/analyze@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 diff --git a/.github/workflows/conformance.yml b/.github/workflows/conformance.yml new file mode 100644 index 0000000..1233deb --- /dev/null +++ b/.github/workflows/conformance.yml @@ -0,0 +1,149 @@ +name: Conformance + +on: + push: + branches: [main] + pull_request: + branches: [main] + workflow_dispatch: + +concurrency: + group: conformance-${{ github.ref }} + cancel-in-progress: true + +env: + CONFORMANCE_VERSION: "0.2.0-alpha.10" + +permissions: + contents: read + +jobs: + server: + runs-on: ubuntu-latest + timeout-minutes: 10 + permissions: + contents: read + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + + - name: Build conformance binaries + run: cargo build -p mcp-conformance + + - name: Test conformance server + run: cargo test -p mcp-conformance --bin conformance-server + + - name: Start conformance server + run: | + PORT=8001 ./target/debug/conformance-server & + echo $! > server.pid + for _ in $(seq 1 30); do + if curl -s -o /dev/null http://127.0.0.1:8001/mcp; then + exit 0 + fi + sleep 1 + done + echo "conformance server did not become ready" >&2 + exit 1 + + - name: Run 2025-11-25 server suite + run: | + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" server \ + --url http://127.0.0.1:8001/mcp \ + --suite all \ + --spec-version 2025-11-25 \ + -o conformance-results/2025-11-25 + + - name: Run 2026-07-28 server suite + run: | + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" server \ + --url http://127.0.0.1:8001/mcp \ + --suite all \ + --spec-version 2026-07-28 \ + -o conformance-results/2026-07-28 + + - name: Run extension server scenarios (informational) + run: | + scenarios=( + tasks-lifecycle + tasks-capability-negotiation + tasks-wire-fields + tasks-request-state-removal + tasks-mrtr-input + tasks-request-headers + tasks-dispatch-and-envelope + tasks-status-notifications + tasks-required-task-error + tasks-mrtr-composition + ) + for scenario in "${scenarios[@]}"; do + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" server \ + --url http://127.0.0.1:8001/mcp \ + --scenario "$scenario" \ + --expected-failures conformance/expected-failures-extensions.yaml \ + -o conformance-extension-results + done + + - name: Stop conformance server + if: always() + run: kill "$(cat server.pid)" 2>/dev/null || true + + - name: Upload results + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: conformance-server-results + path: | + conformance-results + conformance-extension-results + + client: + runs-on: ubuntu-latest + timeout-minutes: 10 + permissions: + contents: read + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + + - name: Build conformance binaries + run: cargo build -p mcp-conformance + + - name: Run 2025-11-25 client suite + run: | + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" client \ + --command "$(pwd)/target/debug/conformance-client" \ + --suite all \ + --spec-version 2025-11-25 \ + -o conformance-client-results/2025-11-25 + + - name: Run 2026-07-28 client suite + run: | + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" client \ + --command "$(pwd)/target/debug/conformance-client" \ + --suite all \ + --spec-version 2026-07-28 \ + -o conformance-client-results/2026-07-28 + + - name: Run extension client suite (informational) + run: | + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" client \ + --command "$(pwd)/target/debug/conformance-client" \ + --suite extensions \ + --expected-failures conformance/expected-failures-extensions.yaml \ + -o conformance-client-results/extensions + + - name: Upload results + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: conformance-client-results + path: conformance-client-results diff --git a/.github/workflows/release-plz.yml b/.github/workflows/release-plz.yml index 04f5119..311b74e 100644 --- a/.github/workflows/release-plz.yml +++ b/.github/workflows/release-plz.yml @@ -1,8 +1,6 @@ name: Release-plz -permissions: - pull-requests: write - contents: write +permissions: {} on: push: @@ -18,19 +16,20 @@ jobs: if: ${{ github.repository_owner == 'modelcontextprotocol' }} permissions: contents: write + pull-requests: read steps: - name: Checkout repository - uses: actions/checkout@v7 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable with: toolchain: "1.92" # Using fork until semver_check_features support is released upstream. # See: https://github.com/release-plz/release-plz/pull/2757 - name: Install release-plz from fork - run: cargo install --locked --git https://github.com/DaleSeo/release-plz --branch feat/semver-check-features release-plz + run: cargo install --locked --git https://github.com/DaleSeo/release-plz --rev c2ea743bacf9bbd56e344df68219e60b0ad35d04 release-plz - name: Run release-plz release run: release-plz release env: @@ -51,17 +50,17 @@ jobs: cancel-in-progress: false steps: - name: Checkout repository - uses: actions/checkout@v7 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable with: toolchain: "1.92" # Using fork until semver_check_features support is released upstream. # See: https://github.com/release-plz/release-plz/pull/2757 - name: Install release-plz from fork - run: cargo install --locked --git https://github.com/DaleSeo/release-plz --branch feat/semver-check-features release-plz + run: cargo install --locked --git https://github.com/DaleSeo/release-plz --rev c2ea743bacf9bbd56e344df68219e60b0ad35d04 release-plz - name: Run release-plz release-pr run: release-plz release-pr env: diff --git a/.github/workflows/scorecard.yml b/.github/workflows/scorecard.yml new file mode 100644 index 0000000..453c0c5 --- /dev/null +++ b/.github/workflows/scorecard.yml @@ -0,0 +1,48 @@ +# OpenSSF Scorecard supply-chain security analysis. +name: Scorecard supply-chain security + +on: + branch_protection_rule: + schedule: + - cron: '5 23 * * 5' + push: + branches: [main] + pull_request: + workflow_dispatch: + +permissions: read-all + +jobs: + analysis: + name: Scorecard analysis + runs-on: ubuntu-latest + permissions: + security-events: write + id-token: write + + steps: + - name: Checkout code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + + - name: Run analysis + uses: ossf/scorecard-action@4eaacf0543bb3f2c246792bd56e8cdeffafb205a # v2.4.3 + with: + results_file: results.sarif + results_format: sarif + # Publishing only succeeds on the default branch. PR and manual runs + # still validate the analysis and produce SARIF without publishing. + publish_results: ${{ github.event_name != 'pull_request' && github.ref == format('refs/heads/{0}', github.event.repository.default_branch) }} + + - name: Upload SARIF artifact + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: scorecard-results + path: results.sarif + retention-days: 5 + + - name: Upload to code scanning + uses: github/codeql-action/upload-sarif@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 + with: + sarif_file: results.sarif diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml new file mode 100644 index 0000000..fe1baa9 --- /dev/null +++ b/.github/workflows/stale.yml @@ -0,0 +1,38 @@ +name: Clean Up Stale PRs and Issues +on: + schedule: + - cron: "0 0 * * *" + workflow_dispatch: + +permissions: + contents: read + +jobs: + stale: + name: Clean Up + runs-on: ubuntu-24.04 + permissions: + issues: write + pull-requests: write + steps: + - uses: actions/stale@4391f3da665fdf50b6810c1a66712fb9ba21aa93 # v11 + with: + days-before-stale: 60 + days-before-close: 14 + exempt-issue-labels: keep,security,pinned + exempt-pr-labels: keep,security,pinned + exempt-draft-pr: true + stale-issue-message: > + This issue has had no activity for 60 days and is now marked as stale. + It will be closed in 14 days if there is no further activity. Add the + `keep` label to keep it open. + close-issue-message: > + Closing this issue after 14 days of inactivity since it was marked stale. + Feel free to reopen if it is still relevant. + stale-pr-message: > + This pull request has had no activity for 60 days and is now marked as stale. + It will be closed in 14 days if there is no further activity. Add the + `keep` label to keep it open. + close-pr-message: > + Closing this pull request after 14 days of inactivity since it was marked stale. + Feel free to reopen if you plan to continue the work. diff --git a/.github/workflows/triage.yml b/.github/workflows/triage.yml index d0d9186..4f1b9ab 100644 --- a/.github/workflows/triage.yml +++ b/.github/workflows/triage.yml @@ -39,7 +39,7 @@ jobs: TRIAGE_MODEL: ${{ vars.TRIAGE_MODEL || 'gpt-4o-mini' }} steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install jq run: sudo apt-get install -y jq diff --git a/Cargo.toml b/Cargo.toml index 192022b..0629e82 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,12 +4,13 @@ default-members = ["crates/rmcp", "crates/rmcp-macros"] resolver = "2" [workspace.dependencies] -rmcp = { version = "2.1.0", path = "./crates/rmcp" } -rmcp-macros = { version = "2.1.0", path = "./crates/rmcp-macros" } +rmcp = { version = "3.4.0", path = "./crates/rmcp" } +rmcp-macros = { version = "3.4.0", path = "./crates/rmcp-macros" } [workspace.package] edition = "2024" -version = "2.1.0" +rust-version = "1.88" +version = "3.4.0" authors = ["4t145 "] license = "Apache-2.0" repository = "https://github.com/modelcontextprotocol/rust-sdk/" diff --git a/DEPENDENCY_POLICY.md b/DEPENDENCY_POLICY.md new file mode 100644 index 0000000..d201e99 --- /dev/null +++ b/DEPENDENCY_POLICY.md @@ -0,0 +1,65 @@ +# Dependency Policy + +This document describes how the Rust MCP SDK (`rmcp`, `rmcp-macros`) selects, +updates, and maintains its dependencies. It applies to all crates in the workspace. + +## Goals + +- Keep dependencies current and secure without churn that destabilizes the public API. +- Minimize the dependency footprint we impose on downstream users. +- Make dependency updates reviewable and traceable. + +## Choosing dependencies + +Before adding a new dependency we consider: + +- **Necessity** — whether the functionality is worth an added dependency, or is small + enough to implement directly. +- **Maintenance & trust** — active maintenance, a healthy release history, and a + permissive, compatible license (the workspace is `Apache-2.0`). +- **Footprint** — transitive dependency count and compile-time cost. +- **Optionality** — anything not needed by every user should sit behind a Cargo + feature so downstream builds stay lean. The SDK is heavily feature-gated for this + reason (transports, `auth`, `schemars`, `reqwest`, etc.). + +Version requirements are specified with caret (`^`, the Cargo default) constraints so +compatible upgrades are picked up automatically. + +## Automated updates + +Dependency updates are automated with [Dependabot](https://docs.github.com/en/code-security/dependabot), +configured in [`.github/dependabot.yml`](.github/dependabot.yml): + +| Ecosystem | Cadence | PR label | +| -------------- | ------- | ----------------- | +| Cargo crates | Weekly | `T-dependencies` | +| GitHub Actions | Daily | `T-CI` | + +Dependabot opens a limited number of PRs at a time (currently 3 per ecosystem) to +keep the review queue manageable. Update PRs run the full CI suite (build, format, +clippy, tests, and conformance) and must be green before merge. + +## Review and merge + +- Patch and minor updates that pass CI are reviewed and merged promptly by maintainers. +- Major (breaking) updates are evaluated individually; they may require code changes + and are scheduled deliberately rather than merged automatically. +- Any dependency change that affects the **public API** (for example a type re-exported + from a dependency) is treated as a potential breaking change and follows + [`VERSIONING.md`](VERSIONING.md). + +## Security updates + +- Security fixes take priority over routine updates and are expedited. +- Vulnerabilities in the SDK itself are handled per [`SECURITY.md`](SECURITY.md) via + GitHub's private Security Advisory process. +- Advisories affecting our dependencies (e.g. via the [RustSec Advisory Database](https://rustsec.org/)) + are addressed as soon as a fixed version is available; if a released version of the + SDK is impacted, a patched release is published and, where warranted, the affected + versions are yanked from crates.io. + +## Minimum Supported Rust Version + +Dependency upgrades must not raise the SDK's declared MSRV (`rust-version` in +`Cargo.toml`) without an explicit, documented MSRV bump — which is itself treated as a +breaking change per [`VERSIONING.md`](VERSIONING.md). diff --git a/README.md b/README.md index 68f3343..7b24d6a 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,3 @@ -
-简体中文 -
- # RMCP [![Crates.io Version](https://img.shields.io/crates/v/rmcp)](https://crates.io/crates/rmcp) [![docs.rs](https://img.shields.io/docsrs/rmcp)](https://docs.rs/rmcp/latest/rmcp) @@ -10,14 +6,20 @@ An official Rust Model Context Protocol SDK implementation with tokio async runtime. -> **Migrating to 1.x?** See the [migration guide](https://github.com/modelcontextprotocol/rust-sdk/discussions/716) for breaking changes and upgrade instructions. +> **Migrating to 3.x?** See the [migration guide](https://github.com/modelcontextprotocol/rust-sdk/discussions/969) for breaking changes and upgrade instructions. This repository contains the following crates: - [rmcp](crates/rmcp): The core crate providing the RMCP protocol implementation - see [rmcp](crates/rmcp/README.md) - [rmcp-macros](crates/rmcp-macros): A procedural macro crate for generating RMCP tool implementations - see [rmcp-macros](crates/rmcp-macros/README.md) -For the full MCP specification, see [modelcontextprotocol.io](https://modelcontextprotocol.io/specification/2025-11-25). +This SDK implements the stable MCP **`2026-07-28`** specification while +remaining fully compatible with the **`2025-11-25`** release and earlier +versions. Features introduced in `2026-07-28` — server discovery & negotiation, +transport-neutral subscriptions, long-running tasks, response caching, +multi-round-trip requests, and standard HTTP routing headers — are documented +below. For the full MCP specification, see +[modelcontextprotocol.io](https://modelcontextprotocol.io/specification/2026-07-28). ## Table of Contents @@ -26,12 +28,21 @@ For the full MCP specification, see [modelcontextprotocol.io](https://modelconte - [Resources](#resources) - [Prompts](#prompts) - [Sampling](#sampling) +- [Elicitation](#elicitation) - [Roots](#roots) - [Logging](#logging) - [Completions](#completions) - [Notifications](#notifications) - [Subscriptions](#subscriptions) +- [Multi-Round-Trip Requests](#multi-round-trip-requests) - [Tasks](#tasks-long-running-tool-invocations) +- [Caching](#caching) +- [Standard HTTP Headers](#standard-http-headers) +- [Stateless Streamable HTTP](#stateless-streamable-http) +- [Transports](#transports) +- [Pagination](#pagination) +- [Capability & Protocol Version Negotiation](#capability--protocol-version-negotiation) +- [JSON Schema 2020-12](#json-schema-2020-12) - [Examples](#examples) - [OAuth Support](#oauth-support) - [Related Resources](#related-resources) @@ -42,10 +53,16 @@ For the full MCP specification, see [modelcontextprotocol.io](https://modelconte ### Import the crate -```toml -rmcp = { version = "0.16.0", features = ["server"] } -## or dev channel -rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", branch = "main" } +Add the latest published version with cargo: + +```sh +cargo add rmcp --features server +``` + +Or use the dev channel: + +```sh +cargo add rmcp --features server --git https://github.com/modelcontextprotocol/rust-sdk --branch main ``` ### Third Dependencies @@ -74,6 +91,44 @@ async fn main() -> Result<(), Box> { ``` +### Client lifecycle modes + +`serve()` uses the legacy MCP lifecycle: the client sends `initialize`, receives +the negotiated server information, and then sends `notifications/initialized`. +Use [`ClientServiceExt::serve_with_lifecycle`](crates/rmcp/src/service/client.rs) to +select another lifecycle explicitly: + +```rust, ignore +use rmcp::{ClientLifecycleMode, ClientServiceExt, ProtocolVersion, model::ClientConfig}; + +// Start directly with server/discover and include client metadata on every request. +let client = ClientConfig::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await?; + +// Or probe the discover lifecycle and fall back when a legacy server reports +// that server/discover is not implemented or does not respond within 10 seconds. +let client = ClientConfig::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await?; +``` + +`ClientLifecycleMode::Initialize` is equivalent to the existing `serve()` behavior. +Discover startup does not send `notifications/initialized`; discovery completes +startup, and each subsequent request carries its protocol version, client +information, and capabilities in `_meta`. + ### Build a Server
@@ -135,7 +190,7 @@ let quit_reason = server.cancel().await?; Tools let servers expose callable functions to clients. Each tool has a name, description, and a JSON Schema for its parameters. Clients discover tools via `list_tools` and invoke them via `call_tool`. -**MCP Spec:** [Tools](https://modelcontextprotocol.io/specification/2025-11-25/server/tools) +**MCP Spec:** [Tools](https://modelcontextprotocol.io/specification/2026-07-28/server/tools) ### Server-side @@ -171,6 +226,11 @@ async fn main() -> anyhow::Result<()> { The generated tool `inputSchema` and `outputSchema` are derived from the fields of `T`. The type name and documentation on `T` are ignored; only field names, field types, and field documentation are used. +> **`2026-07-28` (SEP-2106):** `outputSchema` may now be any JSON Schema type +> (not just `object`), and a tool result's `structuredContent` may be any JSON +> value (string, array, number, …) rather than only an object. Existing +> object-typed tools are unaffected. + When you need custom server metadata or multiple capabilities (tools + prompts), use explicit `#[tool_handler]`: ```rust,ignore @@ -199,6 +259,77 @@ impl ServerHandler for Calculator {} See [`crates/rmcp-macros`](crates/rmcp-macros/README.md) for full macro documentation. +#### Tool result content types + +Beyond a plain `String`, tools can return images, audio, embedded resources, and +mixed content. Build a `CallToolResult` from a `Vec`: + +```rust,ignore +use rmcp::model::{CallToolResult, ContentBlock, ResourceContents}; + +#[tool(description = "Render a chart")] +async fn chart(&self) -> Result { + let png_base64 = render_png(); // base64-encoded image bytes + let wav_base64 = render_wav(); // base64-encoded audio bytes + + Ok(CallToolResult::success(vec![ + // Text + ContentBlock::text("Here is your chart:"), + // Image — base64 data + MIME type + ContentBlock::image(png_base64, "image/png"), + // Audio — base64 data + MIME type + ContentBlock::audio(wav_base64, "audio/wav"), + // Embedded resource — inline text (or ResourceContents::blob for binary) + ContentBlock::resource(ResourceContents::text( + "chart source data", + "chart://last/data.csv", + )), + ])) +} +# fn render_png() -> String { String::new() } +# fn render_wav() -> String { String::new() } +``` + +Image and audio data are base64 strings with a MIME type. For embedded +resources, `ResourceContents::text(..)` inlines text and +`ResourceContents::blob(base64, uri)` inlines binary. + +#### Error handling + +Two failure modes, chosen by **whose problem it is**: + +- **Tool-level error** — `Ok(CallToolResult::error(vec![...]))`. The tool ran but + failed in a way the caller should see (no rows matched, upstream 500). The + client renders your `content`, so the message reaches the user. Use this for + almost every "the tool ran and didn't work" case. +- **Protocol error** — `Err(McpError)` with a JSON-RPC code (e.g. + `McpError::invalid_params(..)`). Use this when the server can't route or process + the request at all; clients render these opaquely, so the caller does **not** + see your message. + +```rust,ignore +use rmcp::model::{CallToolResult, ContentBlock}; +use rmcp::ErrorData as McpError; + +#[tool(description = "Look up a record")] +async fn lookup(&self, Parameters(args): Parameters) -> Result { + // Malformed request — the server can't run anything → protocol error. + if args.query.is_empty() { + return Err(McpError::invalid_params("query must be non-empty", None)); + } + + // Tool ran, no result → tool-level error the user should see. + let rows = self.run_query(&args.query).await; + if rows.is_empty() { + return Ok(CallToolResult::error(vec![ContentBlock::text( + format!("no rows matched '{}'", args.query), + )])); + } + + Ok(CallToolResult::success(vec![ContentBlock::text(format_rows(&rows))])) +} +``` + ### Client-side ```rust,ignore @@ -219,7 +350,7 @@ let result = client.call_tool(CallToolRequestParams::new("add")).await?; Resources let servers expose data (files, database records, API responses) that clients can read. Each resource is identified by a URI and returns content as text or binary (base64-encoded) data. Resource templates allow servers to declare URI patterns with dynamic parameters. -**MCP Spec:** [Resources](https://modelcontextprotocol.io/specification/2025-11-25/server/resources) +**MCP Spec:** [Resources](https://modelcontextprotocol.io/specification/2026-07-28/server/resources) ### Server-side @@ -238,8 +369,8 @@ use serde_json::json; struct MyServer; impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_resources() .build(), @@ -273,6 +404,28 @@ impl ServerHandler for MyServer { "memo://insights" => Ok(ReadResourceResult::new(vec![ ResourceContents::text("Analysis results...", &request.uri), ])), + // Binary resource — base64-encode the bytes and return a blob. + "file:///logo.png" => { + use base64::{Engine, prelude::BASE64_STANDARD}; + let bytes = std::fs::read("logo.png").unwrap_or_default(); + let blob = BASE64_STANDARD.encode(bytes); + Ok(ReadResourceResult::new(vec![ + ResourceContents::blob(blob, &request.uri) + .with_mime_type("image/png"), + ])) + } + // Template-expanded URI — the client fills in `{user_id}` from the + // `users://{user_id}/profile` template declared in + // `list_resource_templates`, and the server reads the concrete URI. + uri if uri.starts_with("users://") && uri.ends_with("/profile") => { + let user_id = uri + .trim_start_matches("users://") + .trim_end_matches("/profile"); + Ok(ReadResourceResult::new(vec![ResourceContents::text( + format!(r#"{{"id": "{user_id}", "name": "User {user_id}"}}"#), + uri, + )])) + } _ => Err(McpError::resource_not_found( "resource_not_found", Some(json!({ "uri": request.uri })), @@ -285,8 +438,12 @@ impl ServerHandler for MyServer { _request: Option, _context: RequestContext, ) -> Result { + // Declare a URI template with a `{user_id}` parameter. Clients expand it + // (e.g. `users://42/profile`) and pass the concrete URI to `read_resource`. Ok(ListResourceTemplatesResult { - resource_templates: vec![], + resource_templates: vec![ + ResourceTemplate::new("users://{user_id}/profile", "user-profile"), + ], next_cursor: None, meta: None, }) @@ -307,8 +464,12 @@ let result = client.read_resource( ReadResourceRequestParams::new("file:///config.json"), ).await?; -// List resource templates +// List resource templates, then read a resource through one by expanding its +// parameters into a concrete URI (`users://{user_id}/profile` → `users://42/profile`). let templates = client.list_all_resource_templates().await?; +let profile = client.read_resource( + ReadResourceRequestParams::new("users://42/profile"), +).await?; ``` ### Notifications @@ -354,7 +515,7 @@ impl ClientHandler for MyClient { Prompts are reusable message templates that servers expose to clients. They accept typed arguments and return conversation messages. The `#[prompt]` macro handles argument validation and routing automatically. -**MCP Spec:** [Prompts](https://modelcontextprotocol.io/specification/2025-11-25/server/prompts) +**MCP Spec:** [Prompts](https://modelcontextprotocol.io/specification/2026-07-28/server/prompts) ### Server-side @@ -421,8 +582,8 @@ impl MyServer { #[prompt_handler] impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_prompts().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_prompts().build()) } } ``` @@ -432,6 +593,30 @@ Prompt functions support several return types: - `GetPromptResult` -- messages with an optional description - `Result` -- either of the above, with error handling +#### Image and embedded-resource content + +A `PromptMessage` can also carry an image or embedded resource. Use the +dedicated constructors (image/audio require the `base64` feature): + +```rust,ignore +use rmcp::model::{PromptMessage, Role}; + +// Image content — raw bytes are base64-encoded for you. +let screenshot: &[u8] = load_png(); +let msg = PromptMessage::new_image(Role::User, screenshot, "image/png", None, None); + +// Embedded resource — inline a text resource by URI. Pass `Some(text)` for a +// text resource, or `None` for a blob resource. +let msg = PromptMessage::new_resource( + Role::User, + "file:///spec.md".to_string(), + Some("text/markdown".to_string()), + Some("# Specification\n...".to_string()), + None, None, None, +); +# fn load_png() -> &'static [u8] { &[] } +``` + ### Client-side ```rust @@ -468,7 +653,7 @@ context.peer.notify_prompt_list_changed().await?; Sampling flips the usual direction: the server asks the client to run an LLM completion. The server sends a `create_message` request, the client processes it through its LLM, and returns the result. -**MCP Spec:** [Sampling](https://modelcontextprotocol.io/specification/2025-11-25/client/sampling) +**MCP Spec:** [Sampling](https://modelcontextprotocol.io/specification/2026-07-28/client/sampling) ### Server-side (requesting sampling) @@ -534,13 +719,147 @@ impl ClientHandler for MyClient { --- +## Elicitation + +Elicitation lets a server pause mid-operation to ask the user for input, in one +of two modes: **form mode** (structured fields with a JSON Schema) or **URL +mode** (send the user to a web page and wait for completion). + +**MCP Spec:** [Elicitation](https://modelcontextprotocol.io/specification/2026-07-28/client/elicitation) + +### Server-side (form mode) + +Define a struct deriving `JsonSchema`, mark it `elicit_safe!`, and call +`elicit::()` on the peer. Schema validation, defaults, and enum choices all +come from the type. + +```rust,ignore +use rmcp::{elicit_safe, model::*, service::{RequestContext, RoleServer}}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +#[schemars(description = "User information")] +pub struct UserInfo { + #[schemars(description = "User's name")] + pub name: String, + // Optional field; omitted if the user doesn't provide it. + #[serde(default)] + #[schemars(description = "Preferred greeting")] + pub greeting: Option, +} + +// Whitelist the type for elicitation (schema-validated on both ends). +elicit_safe!(UserInfo); + +#[tool(description = "Greet the user")] +async fn greet(&self, ctx: RequestContext) -> Result { + // Returns Ok(Some(UserInfo)) if the user accepts. + // Decline and cancel are returned as ElicitationError variants. + match ctx.peer.elicit::("Please provide your name").await { + Ok(Some(info)) => Ok(CallToolResult::success(vec![ContentBlock::text( + format!("Hello, {}!", info.name), + )])), + Ok(None) => Ok(CallToolResult::success(vec![ContentBlock::text( + "No name provided.", + )])), + Err(e) => Ok(CallToolResult::error(vec![ContentBlock::text( + format!("Elicitation failed: {e}"), + )])), + } +} +``` + +#### Enum values + +Enum fields become a choice list. `schemars` needs two hints to inline and type +the enum correctly: + +```rust,ignore +#[derive(Debug, Serialize, Deserialize, JsonSchema, Default)] +#[schemars(inline)] // inline the enum into the parent schema +#[schemars(extend("type" = "string"))] // schemars omits `type` for enums; add it +enum Priority { + #[schemars(title = "Low priority")] + #[default] + Low, + #[schemars(title = "High priority")] + High, +} +``` + +See [`examples/servers/src/elicitation_enum_inference.rs`](examples/servers/src/elicitation_enum_inference.rs) +for single-select, multi-select, titled, and defaulted enum forms. + +### Server-side (URL mode) + +For flows a form can't capture (OAuth consent, a payment page), send the user to +a URL. `elicit_url` returns the user's `ElicitationAction` rather than typed data: + +```rust,ignore +use rmcp::model::ElicitationAction; +use url::Url; + +let action = ctx.peer.elicit_url( + "Please complete setup in your browser", + Url::parse("https://example.com/setup").unwrap(), + "setup-123", // a unique elicitation id +).await?; + +match action { + ElicitationAction::Accept => { /* user consented */ } + ElicitationAction::Decline => { /* user declined */ } + ElicitationAction::Cancel => { /* user aborted */ } +} +``` + +### Client-side + +Implement `ClientHandler::create_elicitation()`, matching on the request variant +to handle form vs. URL mode: + +```rust,ignore +use rmcp::{ClientHandler, model::*, service::{RequestContext, RoleClient}}; + +impl ClientHandler for MyClient { + async fn create_elicitation( + &self, + request: ElicitRequestParams, + _context: RequestContext, + ) -> Result { + match request { + ElicitRequestParams::FormElicitationParams { message, .. } => { + // Show `message` + the requested schema, collect input, then: + Ok(ElicitResult { + action: ElicitationAction::Accept, + content: Some(rmcp::object!({ "name": "Ada" })), + meta: None, + }) + } + ElicitRequestParams::UrlElicitationParams { url, .. } => { + // Open `url`, wait for the user, then report the action. + let _ = url; + Ok(ElicitResult { action: ElicitationAction::Accept, content: None, meta: None }) + } + } + } +} +``` + +On completion the client sends a `notifications/elicitation/response` +notification to release the waiting server-side `elicit_url` call. + +**Example:** [`examples/servers/src/elicitation_stdio.rs`](examples/servers/src/elicitation_stdio.rs) (form + URL), [`examples/servers/src/elicitation_enum_inference.rs`](examples/servers/src/elicitation_enum_inference.rs) (enum forms) + +--- + ## Roots > **Deprecated (SEP-2577):** Roots is deprecated and will be removed in a future release. It remains fully functional for now. See [SEP-2577](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2577). Roots tell servers which directories or projects the client is working in. A root is a URI (typically `file://`) pointing to a workspace or repository. Servers can query roots to know where to look for files and how to scope their work. -**MCP Spec:** [Roots](https://modelcontextprotocol.io/specification/2025-11-25/client/roots) +**MCP Spec:** [Roots](https://modelcontextprotocol.io/specification/2026-07-28/client/roots) ### Server-side @@ -605,7 +924,7 @@ client.notify_roots_list_changed().await?; Servers can send structured log messages to clients. The client sets a minimum severity level, and the server sends messages through the peer notification interface. -**MCP Spec:** [Logging](https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/logging) +**MCP Spec:** [Logging](https://modelcontextprotocol.io/specification/2026-07-28/server/utilities/logging) ### Server-side @@ -615,8 +934,8 @@ Enable the logging capability, handle level changes from the client, and send lo use rmcp::{ServerHandler, model::*, service::RequestContext}; impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_logging() .build(), @@ -678,7 +997,7 @@ client.set_level(SetLevelRequestParams::new(LoggingLevel::Warning)).await?; Completions give auto-completion suggestions for prompt or resource template arguments. As a user fills in arguments, the client can ask the server for suggestions based on what's already been entered. -**MCP Spec:** [Completions](https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/completion) +**MCP Spec:** [Completions](https://modelcontextprotocol.io/specification/2026-07-28/server/utilities/completion) ### Server-side @@ -688,8 +1007,8 @@ Enable the completions capability and implement the `complete()` handler. Use `r use rmcp::{ErrorData as McpError, ServerHandler, model::*, service::RequestContext, RoleServer}; impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_completions() .enable_prompts() @@ -703,6 +1022,7 @@ impl ServerHandler for MyServer { _context: RequestContext, ) -> Result { let values = match &request.r#ref { + // Completion for a prompt argument (`ref/prompt`). Reference::Prompt(prompt_ref) if prompt_ref.name == "sql_query" => { match request.argument.name.as_str() { "operation" => vec!["SELECT", "INSERT", "UPDATE", "DELETE"], @@ -723,6 +1043,17 @@ impl ServerHandler for MyServer { _ => vec![], } } + // Completion for a resource-template argument (`ref/resource`). The + // `uri` identifies the template (e.g. `users://{user_id}/profile`) + // and `argument.name` is the template variable being completed. + Reference::Resource(resource_ref) + if resource_ref.uri == "users://{user_id}/profile" => + { + match request.argument.name.as_str() { + "user_id" => vec!["1", "2", "42"], + _ => vec![], + } + } _ => vec![], }; @@ -744,12 +1075,22 @@ impl ServerHandler for MyServer { ```rust use rmcp::model::*; +// Completion for a prompt argument. let result = client.complete(CompleteRequestParams::new( Reference::for_prompt("sql_query"), ArgumentInfo::new("operation", "SEL"), )).await?; // result.completion.values contains suggestions like ["SELECT"] + +// Completion for a resource-template argument: reference the template by URI +// and complete one of its variables (`user_id`). +let resource_completion = client.complete(CompleteRequestParams::new( + Reference::for_resource("users://{user_id}/profile"), + ArgumentInfo::new("user_id", "4"), +)).await?; + +// resource_completion.completion.values contains suggestions like ["42"] ``` **Example:** [`examples/servers/src/completion_stdio.rs`](examples/servers/src/completion_stdio.rs) @@ -760,7 +1101,7 @@ let result = client.complete(CompleteRequestParams::new( Notifications are fire-and-forget messages -- no response is expected. They cover progress updates, cancellation, and lifecycle events. Both sides can send and receive them. -**MCP Spec:** [Notifications](https://modelcontextprotocol.io/specification/2025-11-25/basic/notifications) +**MCP Spec:** [Notifications](https://modelcontextprotocol.io/specification/2026-07-28/basic#notifications) ### Progress notifications @@ -810,12 +1151,62 @@ impl ServerHandler for MyServer { } ``` -### Initialized notification +### Ping -Clients send `initialized` after the handshake completes: +Either side can send a `ping` request to check that its counterpart is still +responsive and the connection is alive. A ping carries no parameters and the +receiver replies with an empty result. Because pings can flow in both +directions, `rmcp` handles them symmetrically: + +- **Sending a ping** — construct a `PingRequest` and send it over the peer. + A client pings the server with `ClientRequest::PingRequest`; a server pings + the client with `ServerRequest::PingRequest`. `send_request` resolves once the + empty response arrives, so a returned `Ok` confirms the peer is reachable: ```rust -// Sent automatically by rmcp during the serve() handshake. +use rmcp::model::{PingRequest, ServerRequest}; + +// From a server, ping the connected client to verify it is still alive. +context.peer + .send_request(ServerRequest::PingRequest(PingRequest::default())) + .await?; +``` + +```rust +use rmcp::model::{ClientRequest, PingRequest}; + +// From a client, ping the server. `running` is the value returned by serve(). +running + .send_request(ClientRequest::PingRequest(PingRequest::default())) + .await?; +``` + +- **Responding to a ping** — `rmcp` answers incoming pings automatically. The + default `ping` method on `ServerHandler` and `ClientHandler` returns an empty + result, so no code is required. Override it only if you want to run custom + logic (for example, health checks) when a ping arrives: + +```rust +impl ServerHandler for MyServer { + async fn ping( + &self, + _context: RequestContext, + ) -> Result<(), McpError> { + // Custom liveness logic here, if any. + Ok(()) + } +} +``` + +**MCP Spec:** [Ping](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/ping) + +### Initialized notification + +Legacy clients send `initialized` after the `initialize` handshake completes. +Clients using `ClientLifecycleMode::Discover` do not send this notification: + +```rust +// Sent automatically by rmcp during the legacy serve() handshake. // Servers handle it via: impl ServerHandler for MyServer { async fn on_initialized( @@ -843,114 +1234,581 @@ context.peer.notify_resource_list_changed().await?; ## Subscriptions -Clients can subscribe to specific resources. When a subscribed resource changes, the server sends a notification and the client can re-read it. +Protocol `2026-07-28` replaces `resources/subscribe`, `resources/unsubscribe`, and +the standalone HTTP GET stream with the transport-neutral, long-lived +`subscriptions/listen` request. Each requested notification category is opt-in. -**MCP Spec:** [Resources - Subscriptions](https://modelcontextprotocol.io/specification/2025-11-25/server/resources#subscriptions) +**MCP Spec:** [Subscriptions](https://modelcontextprotocol.io/specification/2026-07-28/basic/patterns/subscriptions) ### Server-side -Enable subscriptions in the resources capability and implement the `subscribe()` / `unsubscribe()` handlers: +Declare the notification capabilities you serve, return the accepted subset, +and use the filter-enforcing subscription sink: ```rust -use rmcp::{ErrorData as McpError, ServerHandler, model::*, service::RequestContext, RoleServer}; -use std::sync::Arc; -use tokio::sync::Mutex; -use std::collections::HashSet; - -#[derive(Clone)] -struct MyServer { - subscriptions: Arc>>, -} +use rmcp::{ + ErrorData, ServerHandler, + model::*, + service::SubscriptionContext, +}; impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() - .enable_resources() - .enable_resources_subscribe() + .enable_tools() + .enable_tool_list_changed() .build(), ) } - async fn subscribe( + fn accepted_subscription_filter( &self, - request: SubscribeRequestParams, - _context: RequestContext, - ) -> Result<(), McpError> { - self.subscriptions.lock().await.insert(request.uri); - Ok(()) + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.clone()) } - async fn unsubscribe( - &self, - request: UnsubscribeRequestParams, - _context: RequestContext, - ) -> Result<(), McpError> { - self.subscriptions.lock().await.remove(&request.uri); + async fn listen(&self, context: SubscriptionContext) -> Result<(), ErrorData> { + if context.accepted().tools_list_changed == Some(true) { + context.sink().notify_tool_list_changed().await + .map_err(|error| ErrorData::internal_error(error.to_string(), None))?; + } + context.cancelled().await; Ok(()) } } ``` -When a subscribed resource changes, notify the client: - -```rust -// Check if the resource has subscribers, then notify -context.peer.notify_resource_updated( - ResourceUpdatedNotificationParam::new("file:///config.json"), -).await?; -``` +The SDK intersects the handler's filter with the requested categories and the +capabilities advertised by `get_info()`. It sends the acknowledgment before +`listen`, tags every sink notification with the listen request ID, and rejects +categories or resource URIs outside the accepted filter. ### Client-side ```rust use rmcp::model::*; -// Subscribe to updates for a resource -client.subscribe(SubscribeRequestParams::new("file:///config.json")).await?; +let mut subscription = client.listen( + SubscriptionFilter::builder() + .tools_list_changed() + .resource_subscription("file:///config.json") + .build(), +).await?; -// Unsubscribe when no longer needed -client.unsubscribe(UnsubscribeRequestParams::new("file:///config.json")).await?; +println!("accepted: {:?}", subscription.acknowledged()); +while let Some(notification) = subscription.next().await? { + println!("notification: {notification:?}"); +} + +subscription.cancel().await?; ``` -Handle update notifications in `ClientHandler`: +`listen()` buffers up to 64 notifications per subscription. Use +`listen_with_capacity()` to choose a different non-zero capacity; if a consumer +falls behind, `Subscription::end()` reports `SubscriptionEnd::Lagged`. -```rust -impl ClientHandler for MyClient { - async fn on_resource_updated( - &self, - params: ResourceUpdatedNotificationParam, - _context: NotificationContext, - ) { - // Re-read the resource at params.uri +For older negotiated protocol versions, the deprecated `subscribe()` and +`unsubscribe()` APIs retain their legacy wire behavior. Modern Streamable HTTP +uses the listen POST response stream directly and does not use sessions, GET, +DELETE, or `Last-Event-ID`. After an abrupt transport close, call `listen` +again; subscription state is not resumed across HTTP or stdio reconnects. + +See the +[modern subscription server](examples/servers/src/subscriptions_streamhttp.rs) +and [client](examples/clients/src/subscriptions_streamhttp.rs) examples. + +--- + +## Multi-Round-Trip Requests + +Protocol `2026-07-28` adds Multi-Round-Trip Requests (MRTR, SEP-2322): a server +can answer a `tools/call`, `prompts/get`, or `resources/read` with an +`InputRequiredResult` instead of a final result, asking the client to fulfill +one or more embedded server requests (elicitation, sampling, or roots) and then +retry. The exchange is stateless — the server carries its progress in an opaque +`requestState` that the client echoes back verbatim. + +**MCP Spec:** [Multiple Round-Trip Requests](https://modelcontextprotocol.io/specification/2026-07-28/server/tools#multiple-round-trip-requests) + +### Server-side + +Return an `InputRequiredResult` via the outcome enum for the method +(`CallToolResponse`, `GetPromptResponse`, or `ReadResourceResponse`). The SDK +only forwards it to peers that negotiated `2026-07-28` or newer — older peers +get a protocol error instead. + +```rust, ignore +async fn call_tool(&self, request: CallToolRequestParams, _ctx: RequestContext) + -> Result +{ + match request.request_state { + // First round: ask the client for input, seal progress into requestState. + None => { + let mut input_requests = InputRequests::new(); + input_requests.insert("city".into(), InputRequest::Elicitation(elicit_city())); + let sealed = self.codec.seal_json(&json!({ "awaiting": "city" }))?; + Ok(InputRequiredResult::new(Some(input_requests), Some(sealed)).into()) + } + // Retry round: verify the echoed state, read the responses, finish. + Some(sealed) => { + let _state = self.codec.open_json(&sealed) + .map_err(|_| ErrorData::invalid_params("tampered request state", None))?; + let city = request.input_responses.as_ref() + .and_then(|r| r.get("city")); + Ok(CallToolResult::success(vec![ContentBlock::text("It is sunny.")]).into()) + } } } ``` +> **`requestState` is untrusted.** [SEP-2322 requires servers to validate +> it](https://modelcontextprotocol.io/seps/2322-MRTR#protocol-requirements-for-ephemeral-workflow) +> because the client echoes it back verbatim. A stateless server that stores +> meaningful data in it MUST verify integrity first. Enable the `request-state` +> feature and use `RequestStateCodec` to seal and open it (HMAC-tagged), or keep +> state server-side and use `requestState` only as an opaque handle. + +For multi-replica deployments, use `RequestStateCodec::new_with_keyring` to +rotate signing keys without invalidating in-flight requests: + +1. Deploy the old and new keys everywhere, continuing to emit `rs1` with the + old key via `with_rs1_signing("old")`. +2. Start emitting `rs2` with the new key while retaining the old key via + `with_rs1_fallback("old")`. +3. After the maximum `requestState` lifetime has elapsed, remove the old key. + +### Client-side + +The high-level `call_tool`, `get_prompt`, and `read_resource` helpers drive MRTR +automatically: they fulfill each embedded request through the local +`ClientHandler` and retry, up to `DEFAULT_MRTR_MAX_ROUNDS` (10). + +```rust, ignore +// Auto mode: the SDK fulfills embedded requests and retries for you. +let result = client.call_tool(CallToolRequestParams::new("weather")).await?; + +// Choose a custom round cap. +let result = client + .call_tool_with_mrtr_max_rounds(CallToolRequestParams::new("weather"), 3) + .await?; + +// Manual mode: get the intermediate InputRequiredResult and drive rounds yourself. +match client.call_tool_once(CallToolRequestParams::new("weather")).await? { + CallToolResponse::InputRequired(input_required) => { /* fulfill + retry */ } + CallToolResponse::Complete(result) => { /* done */ } + _ => {} +} +``` + +**Example:** [`examples/servers/src/mrtr.rs`](examples/servers/src/mrtr.rs) (end-to-end server + client) + --- ## Tasks (long-running tool invocations) -`rmcp` supports the [task-based tool invocation](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks) -flow defined in SEP-1319. Annotate a tool with `execution(task_support = "required" | "optional")` -and add `#[task_handler]` to your `ServerHandler` impl — `enqueue_task`, `tasks/list`, `tasks/get`, -`tasks/result`, and `tasks/cancel` are generated for you on top of an `OperationProcessor`. +`rmcp` implements the [MCP Tasks extension](https://modelcontextprotocol.io/extensions/tasks/overview) +(SEP-2663, `io.modelcontextprotocol/tasks`). A client declares the extension in its +capabilities; the server then decides per request whether to materialize a `tools/call` +as a task, returning a `CreateTaskResult` (`resultType: "task"`). The client polls +`tasks/get`, answers in-task input requests via `tasks/update`, and may request +cooperative cancellation via `tasks/cancel`. Use `rmcp::task_manager::TaskManager` +to manage task lifecycles server-side. ```rust, ignore -#[tool( - description = "Sum two numbers after a 2-second delay", - execution(task_support = "required") -)] -async fn slow_sum(/* ... */) -> Result { /* ... */ } +// Client: declare the tasks extension capability. +let caps = ClientCapabilities::builder().enable_tasks().build(); -#[tool_handler] -#[task_handler] -impl ServerHandler for TaskDemo {} +// Server: decide per request whether to materialize a task. +async fn call_tool(&self, request: CallToolRequestParams, context: RequestContext) + -> Result +{ + let client_supports_tasks = context + .client_capabilities() + .is_some_and(|caps| caps.supports_tasks()); + if client_supports_tasks { + let task = self.tasks.spawn(TaskOptions::default(), move |_ctx| { + Box::pin(async move { /* long-running work -> Ok(CallToolResult) */ }) + }); + return Ok(CallToolResponse::Task(CreateTaskResult::new(task))); + } + // ... fall back to synchronous execution +} ``` See [`servers_task_stdio`](examples/servers/src/task_stdio.rs) and the matching [`clients_task_stdio`](examples/clients/src/task_stdio.rs) for a runnable end-to-end example. +## Caching + +`rmcp` clients transparently cache responses that carry the +[SEP-2549](https://modelcontextprotocol.io/specification/2026-07-28/server/utilities/caching) +caching hints (`ttlMs` / `cacheScope`) for `server/discover`, `tools/list`, +`prompts/list`, `resources/list`, `resources/templates/list`, and `resources/read`. + +Caching is on by default but only stores a response when the server sends a +positive `ttlMs`, so servers that omit the hint behave exactly as before. Entries +expire after their TTL, are partitioned by cache scope, and are invalidated +automatically by the matching `list_changed` / `resource updated` notifications. + +No call-site changes are needed — existing calls benefit automatically: + +```rust, ignore +let tools = peer.list_tools(None).await?; // served from cache while fresh +let res = peer.read_resource(params).await?; // cached per-URI +``` + +Tune or disable it per connection via the `Peer`: + +```rust, ignore +use std::time::Duration; +use rmcp::ClientCacheConfig; + +// Customize behavior. +peer.set_response_cache_config( + ClientCacheConfig::default() + .with_default_ttl(Duration::from_secs(30)) // TTL for servers that omit ttlMs + .with_max_ttl(Duration::from_secs(3600)) // upper bound on any TTL + .with_max_entries(1024) + .with_private_partition(user_id) // separate private caches per principal + .with_serve_stale_on_error(false), // surface errors instead of stale data +).await; + +// Or turn it off entirely. +peer.set_response_cache_config(ClientCacheConfig::disabled()).await; + +// Manually flush. +peer.clear_response_cache().await; +``` + +> **Note:** with the default `serve_stale_on_error`, a failed re-fetch returns the +> last cached response (even if expired) as `Ok(..)` instead of an error. Set +> `with_serve_stale_on_error(false)` if callers must observe fetch failures. + +## Standard HTTP Headers + +Protocol `2026-07-28` standardizes a set of Streamable HTTP request headers +(SEP-2243) so proxies and gateways can route MCP traffic without parsing the +JSON body: `Mcp-Method`, `Mcp-Name`, and `Mcp-Param-*`. `rmcp` emits and +validates these automatically once a connection negotiates `2026-07-28` or +newer — no call-site changes are required, and older negotiated versions are +untouched. + +**MCP Spec:** [Header standardization](https://modelcontextprotocol.io/specification/2026-07-28/basic/transports#header) + +- `Mcp-Method` — the JSON-RPC method (e.g. `tools/call`). +- `Mcp-Name` — the target name, sourced from `params.name` (`tools/call`, + `prompts/get`), `params.uri` (`resources/*`), or `params.taskId` (`tasks/*`). +- `Mcp-Param-*` — selected `tools/call` arguments, promoted from the tool's + input schema. + +To promote a tool argument into a routing header, annotate the top-level schema +property with `x-mcp-header`: + +```rust, ignore +// A `region` argument surfaces as the `Mcp-Param-Region` request header. +let schema = serde_json::json!({ + "type": "object", + "properties": { + "region": { "type": "string", "x-mcp-header": "Region" } + } +}); +``` + +Annotations must be non-empty RFC 9110 tokens, case-insensitively unique, and +applied only to top-level primitive (`string`/`integer`/`boolean`) properties. +Values that cannot travel as a bare header (leading/trailing whitespace, +control/non-ASCII characters) are transparently Base64-wrapped as +`=?base64??=`. + +--- + +## Stateless Streamable HTTP + +Per SEP-2567, `rmcp` serves the `2026-07-28` draft statelessly **automatically**: +no `Mcp-Session-Id`, no standalone GET/DELETE stream, and no `Last-Event-ID` +resumption. The `legacy_session_mode` flag below only controls behavior for +*legacy* protocol versions (`< 2026-07-28`). + +**MCP Spec:** [Transports](https://modelcontextprotocol.io/specification/2026-07-28/basic/transports) + +### Server-side + +A default server already serves `2026-07-28` clients statelessly. To also serve +*legacy* clients without sessions, disable `legacy_session_mode` (formerly +`stateful_mode`; builder `with_stateful_mode`). Optionally set +`with_json_response(true)` so simple request/response tools reply with a single +`application/json` body instead of an SSE stream (the server still falls back to +`text/event-stream` if a handler emits a notification or server request first): + +```rust, ignore +use rmcp::transport::streamable_http_server::{ + StreamableHttpService, StreamableHttpServerConfig, + session::local::LocalSessionManager, +}; + +let config = StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) // stateless for legacy versions too + .with_json_response(true); // plain JSON replies for simple tools + +let service = StreamableHttpService::new( + || Ok(Counter::new()), // a fresh handler per request + LocalSessionManager::default().into(), + config, +); + +// `StreamableHttpService` is a Tower service — mount it on any router. +let router = axum::Router::new().nest_service("/mcp", service); +``` + +> Because there is no per-session state, the `service_factory` runs per request. +> Keep shared state (DB pools, caches) in a `Clone` handle captured by the +> closure; don't rely on in-memory state surviving between requests. +> +> Modern-only servers can additionally call +> `with_stateless_protocol_metadata_required(true)` to reject the compatibility +> fallback for requests missing their per-request protocol signals. rmcp clients +> negotiated below `2026-07-28` do not attach that body metadata and will be +> rejected, so pair this option with a `supported_protocol_versions` +> implementation that advertises only `2026-07-28` and later. + +### Client-side + +The Streamable HTTP client transport allows stateless operation by default +(`allow_stateless: true`), so no configuration is needed to talk to a stateless +server — it simply omits the session header when the server doesn't issue one: + +```rust, ignore +use rmcp::transport::StreamableHttpClientTransport; + +// Defaults are stateless-friendly. +let transport = StreamableHttpClientTransport::from_uri("http://localhost:8000/mcp"); +let client = ClientConfig::default().serve(transport).await?; +``` + +**Example:** [`examples/servers/src/counter_streamhttp.rs`](examples/servers/src/counter_streamhttp.rs) (server), [`examples/clients/src/streamable_http.rs`](examples/clients/src/streamable_http.rs) (client) + +--- + +## Transports + +A transport moves JSON-RPC messages between client and server. Any `Transport` +impl can be passed to `.serve(..)`; `rmcp` ships the common ones behind Cargo +features. + +**MCP Spec:** [Transports](https://modelcontextprotocol.io/specification/2026-07-28/basic/transports) + +| Transport | Feature(s) | Notes | +| ---------------------------- | -------------------------------------------- | ----- | +| **stdio** | `transport-io` (client + server) | Communicate over `stdin`/`stdout`; the standard way to launch local MCP servers as child processes. | +| **Child process** (client) | `transport-child-process` | Spawn a server binary and talk to it over its stdio. | +| **Streamable HTTP** (server) | `transport-streamable-http-server` | The current HTTP transport. Exposes a Tower service you can mount on any router. | +| **Streamable HTTP** (client) | `transport-streamable-http-client-reqwest` | HTTP client transport built on `reqwest`. | +| **Worker / in-process** | `transport-worker` | For embedding or testing without real I/O. | + +### stdio + +```rust,ignore +use rmcp::{ServiceExt, transport::stdio}; + +// Server: serve over stdin/stdout. +let server = MyServer.serve(stdio()).await?; +server.waiting().await?; +``` + +```rust,ignore +use rmcp::{ServiceExt, transport::{TokioChildProcess, ConfigureCommandExt}}; +use tokio::process::Command; + +// Client: launch a server binary and talk to it over its stdio. +let transport = TokioChildProcess::new(Command::new("uvx").configure(|cmd| { + cmd.arg("mcp-server-git"); +}))?; +let client = ().serve(transport).await?; +``` + +### Streamable HTTP + +`StreamableHttpService` is a Tower service — mount it on any `axum`/`hyper` +router (see [Stateless Streamable HTTP](#stateless-streamable-http) for the full +server example). The client transport connects with a single URI: + +```rust,ignore +use rmcp::transport::StreamableHttpClientTransport; + +let transport = StreamableHttpClientTransport::from_uri("http://localhost:8000/mcp"); +let client = ClientConfig::default().serve(transport).await?; +``` + +The client allows up to 16 ordinary http POSTs at once. Configure this with +`StreamableHttpClientTransportConfig::with_uri(url).max_concurrent_requests(n)`; +`1` keeps ordinary POSTs serial, and `0` is treated as `1`. An open sse response +stream does not count against this limit. Cancellation and replies use a +separate queue with one extra POST slot. Configure their timeout with +`control_request_timeout` (default: five seconds). The timeout starts when the +POST starts, excluding time in the queue. Cancellation stops a queued or active +POST immediately. +For an open legacy response stream, the client stops reading but keeps the stream +alive until the cancellation send finishes or is dropped. This lets custom http +adapters handle cancellation before their stream state is removed. + +Session recovery waits up to five seconds for old POSTs, then stops any that +remain. Those POSTs are not retried because the server may have processed them. +Configure this wait and the separate +reinitialization timeout with `session_recovery_timeout`. Callers still decide +which tools may run at the same time and which need approval. + +#### Server-Sent Events (SSE) + +Streamable HTTP responses arrive as either a single `application/json` body or a +`text/event-stream` (Server-Sent Events) stream when the server pushes +notifications or requests before the result. `rmcp` handles both automatically +(SSE parsing lives behind the `client-side-sse` feature). There is no separate +"SSE transport" to configure — it's an implementation detail of Streamable HTTP. + +#### Legacy HTTP+SSE transport (`2024-11-05`) — intentionally not provided + +The standalone two-endpoint **HTTP+SSE transport** defined in protocol revision +`2024-11-05` (a separate `GET` SSE channel plus a `POST` message endpoint) is a +**deliberate non-goal** for `rmcp`. It was [replaced by Streamable HTTP in the +`2025-03-26` revision](https://modelcontextprotocol.io/specification/2026-07-28/basic/transports), +and `rmcp` targets current spec revisions (`2025-11-25` and `2026-07-28`), so it +ships **no legacy HTTP+SSE client or server transport**. + +What to use instead: + +- **New client/server code** — use [Streamable HTTP](#streamable-http). It carries + the same SSE streaming semantics over a single endpoint and is the transport all + supported spec revisions expect. +- **Server-to-client streaming** (push notifications, resource updates) — this is + built into Streamable HTTP; see [Subscriptions](#subscriptions). +- **Talking to a legacy `2024-11-05`-only server** — front it with a proxy that + speaks Streamable HTTP, or pin a dependency to a release that predates the + transport's removal. `rmcp` will not add the legacy transport back. + +This is a supported-surface decision, not a missing feature: every transport +`rmcp` implements is listed in the [Transports](#transports) table above. + +--- + +## Pagination + +List operations (`tools/list`, `prompts/list`, `resources/list`, +`resources/templates/list`) are paginated via a `next_cursor`. The `list_all_*` +helpers walk every page for you: + +```rust,ignore +// Fetches all pages transparently. +let tools = client.list_all_tools().await?; +let prompts = client.list_all_prompts().await?; +let resources = client.list_all_resources().await?; +``` + +To page manually, call the single-page method and follow `next_cursor` until +it's `None`: + +```rust,ignore +use rmcp::model::PaginatedRequestParams; + +let mut cursor = None; +loop { + let page = client + .list_tools(Some(PaginatedRequestParams { meta: None, cursor })) + .await?; + for tool in &page.tools { + // handle each tool + } + cursor = page.next_cursor; + if cursor.is_none() { + break; + } +} +``` + +On the server, return a `next_cursor` from your `list_*` handler when more pages +remain (`None` when complete). + +**MCP Spec:** [Pagination](https://modelcontextprotocol.io/specification/2026-07-28/server/utilities/pagination) + +--- + +## Capability & Protocol Version Negotiation + +### Capabilities + +At initialization, client and server exchange **capabilities** so each side +knows what the other supports. Declare yours with the `ServerCapabilities` +builder in `get_info()`: + +```rust,ignore +use rmcp::model::{ServerCapabilities, ServerConfig}; + +fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_prompts() + .enable_resources() + .enable_resources_subscribe() + .enable_tool_list_changed() + .enable_logging() + .build(), + ) +} +``` + +Clients do the same via `ClientCapabilities::builder()`. Macros like +`#[tool_handler]` / `#[prompt_handler]` set the relevant flags automatically. +After connecting, read the peer's capabilities via `peer.peer_info()`. + +### Protocol version + +MCP is versioned by date. `rmcp` negotiates automatically on connect — the +client offers a preferred `ProtocolVersion` and falls back to one the server +supports: + +```rust,ignore +use rmcp::model::ProtocolVersion; + +ProtocolVersion::LATEST; // newest stable version this SDK defaults to +ProtocolVersion::V_2026_07_28; // a specific version constant +ProtocolVersion::KNOWN_VERSIONS; // every version this SDK understands +``` + +Version-specific behavior (SEP-2243 headers, SEP-2567 stateless serving, the +SEP-2575 subscription model) is gated on the negotiated version, so older clients +keep working while newer ones opt in. + +**MCP Spec:** [Versioning and Compatibility](https://modelcontextprotocol.io/specification/2026-07-28/basic/versioning) + +--- + +## JSON Schema 2020-12 + +Deriving `schemars::JsonSchema` on your parameter and result types generates +[JSON Schema draft 2020-12](https://json-schema.org/) — the dialect the MCP spec +requires — for the tool's `inputSchema` and `outputSchema`. + +```rust,ignore +use rmcp::schemars; + +#[derive(Debug, serde::Deserialize, schemars::JsonSchema)] +struct SearchParams { + /// Full-text query. + query: String, + /// Maximum number of results. + #[serde(default)] + limit: Option, +} +``` + +Field names, types, and doc comments flow into the schema — no manual authoring +needed. As of `2026-07-28` (SEP-2106), `outputSchema` may be any JSON Schema type +(not only `object`) and `structuredContent` may be any JSON value. + +--- + ## Examples See [examples](examples/README.md). @@ -961,8 +1819,8 @@ See [Oauth_support](docs/OAUTH_SUPPORT.md) for details. ## Related Resources -- [MCP Specification](https://modelcontextprotocol.io/specification/2025-11-25) -- [Schema](https://github.com/modelcontextprotocol/specification/blob/main/schema/2025-11-25/schema.ts) +- [MCP Specification](https://modelcontextprotocol.io/specification/2026-07-28) +- [Schema](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/2026-07-28/schema.ts) ## Related Projects diff --git a/ROADMAP.md b/ROADMAP.md index 4910bec..6eea344 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -1,71 +1,112 @@ # RMCP Roadmap -This roadmap tracks the path to SEP-1730 Tier 1 for the Rust MCP SDK. +This roadmap tracks the path to [SEP-1730](https://github.com/modelcontextprotocol/modelcontextprotocol/issues/1730) Tier 1 for the Rust MCP SDK. -Server conformance: 87.5% (28/32) · Client conformance: 80.0% (16/20) +**Status: all SEP-1730 Tier 1 requirements are met.** Conformance is +100% across every date-versioned suite, a stable ≥1.0.0 release has shipped +([Releases](https://github.com/modelcontextprotocol/rust-sdk/releases)), +issue triage and critical-bug (P0) resolution are within the Tier 1 SLAs, the +governance documents (`VERSIONING.md`, `DEPENDENCY_POLICY.md`, this `ROADMAP.md`) are +published, and all 48 non-experimental features are documented with examples. This +document now serves as the ongoing tracker for spec conformance and SDK health. + +| SEP-1730 Tier 1 requirement | Status | Evidence | +| --------------------------------------------------- | ------ | -------- | +| Server conformance 100% (date-versioned) | ✅ | 30/30 — see below | +| Client conformance 100% (date-versioned) | ✅ | 39/39 scored — see below | +| Issue triage ≥90% within 2 business days | ✅ | 95.2% (20/21) | +| All P0 bugs resolved within 7 days | ✅ | 0 open; last P0 (#741) resolved in 3 days | +| Stable release ≥1.0.0 (no pre-release suffix) | ✅ | [latest release](https://github.com/modelcontextprotocol/rust-sdk/releases/latest) (see tooling note below) | +| Clear versioning + breaking-change policy | ✅ | [`VERSIONING.md`](VERSIONING.md) | +| All non-experimental features documented w/ examples| ✅ | 48/48 in [`README.md`](README.md) | +| Published dependency update policy | ✅ | [`DEPENDENCY_POLICY.md`](DEPENDENCY_POLICY.md) + [`.github/dependabot.yml`](.github/dependabot.yml) | +| Published roadmap tracking spec components | ✅ | this document | + +| Suite (date-versioned) | Server | Client | +| ---------------------- | ------------- | ------------- | +| 2025-11-25 | 100% (30/30) | 100% | +| 2026-07-28 | 100% (30/30) | 100% | + +Only date-versioned scenarios count toward SDK tiering. `draft` (2026-07-28 draft) +and `extension` scenarios are informational and reported separately below. + +> **Tooling note — `stable_release`:** the SEP-1730 `tier-check` CLI may report +> `stable_release` as failing because it does not parse the workspace tag prefix +> `rmcp-v` (as in `rmcp-vX.Y.Z`). The published releases are genuine stable, +> non-pre-release versions ≥1.0.0 +> ([Releases](https://github.com/modelcontextprotocol/rust-sdk/releases)); the +> flag is a tooling artifact, not an unmet requirement. --- -## Tier 2 → Tier 1 +## Conformance -### Conformance +Per-scenario status for the current spec is tracked in the epic issue +[#977 — Tracking: 2026-07-28 spec conformance](https://github.com/modelcontextprotocol/rust-sdk/issues/977), +under the [`2026-07-28 spec` milestone](https://github.com/modelcontextprotocol/rust-sdk/milestone/3). -#### Server (87.5% → 100%) +CI (`.github/workflows/conformance.yml`) runs the full `2025-11-25` and `2026-07-28` +server and client suites on every push and PR. Both suites are fully green. -- [ ] Fix `prompts-get-with-args` — prompt argument handling returns incorrect result (arg1/arg2 not substituted) -- [ ] Fix `prompts-get-embedded-resource` — embedded resource content in prompt responses (invalid content union) -- [ ] Fix `elicitation-sep1330-enums` — enum inference handling per SEP-1330 (missing enumNames for legacy titled enum) -- [ ] Fix `dns-rebinding-protection` — validate `Host` / `Origin` headers on Streamable HTTP transport (accepts invalid headers with 200) +### Informational (not scored for tiering) -#### Client (80.0% → 100%) +Extension-tagged scenarios are excluded by `--spec-version` filters, so CI runs them +in separate steps against `conformance/expected-failures-extensions.yaml`: -- [ ] Fix `auth/metadata-var3` — AS metadata discovery variant 3 (no authorization support detected) -- [ ] Fix `auth/scope-from-www-authenticate` — use scope parameter from WWW-Authenticate header on 403 insufficient_scope -- [ ] Fix `auth/scope-step-up` — handle 403 `insufficient_scope` and re-authorize with upgraded scopes -- [ ] Fix `auth/2025-03-26-oauth-endpoint-fallback` — legacy OAuth endpoint fallback for pre-2025-06-18 servers (no authorization support detected) +| Scenario | Tag | Status | +| --------------------------------------- | ---------- | ------ | +| `auth/client-credentials-basic` | extension | ✅ Pass | +| `auth/client-credentials-jwt` | extension | ✅ Pass | +| `auth/enterprise-managed-authorization` | extension | ❌ Expected failure (not implemented by the conformance client) | +| `auth/wif-jwt-bearer` | 2026-07-28 draft | ❌ Expected failure (WIF / SEP-1933, draft) | +| `tasks-*` (SEP-2663) | extension | ❌ 9 expected failures · ⏭️ 1 upstream-skipped | -### Governance & Policy +### Spec features without conformance scenarios -- [ ] Create `VERSIONING.md` — document semver scheme, what constitutes a breaking change, and how breaking changes are communicated +Conformance does not cover the entire spec surface. Remaining feature work tracked via +the milestone: -### Documentation (26/48 → 48/48 features with prose + examples) - -#### Undocumented features (14) - -- [ ] Tools — image results -- [ ] Tools — audio results -- [ ] Tools — embedded resources -- [ ] Prompts — embedded resources -- [ ] Prompts — image content -- [ ] Elicitation — URL mode -- [ ] Elicitation — default values -- [ ] Elicitation — complete notification -- [ ] Ping -- [ ] SSE transport — legacy (client) -- [ ] SSE transport — legacy (server) -- [ ] Pagination -- [ ] Protocol version negotiation -- [ ] JSON Schema 2020-12 support *(upgrade from partial)* - -#### Partially documented features (7) - -- [ ] Tools — error handling *(add dedicated prose + example)* -- [ ] Resources — reading binary *(add dedicated example)* -- [ ] Elicitation — form mode *(add prose docs, not just example README)* -- [ ] Elicitation — schema validation *(add prose docs)* -- [ ] Elicitation — enum values *(add prose docs)* -- [ ] Capability negotiation *(add dedicated prose explaining the builder API)* -- [ ] Protocol version negotiation *(document version negotiation behavior)* +- SEP-2567 sessionless MCP via explicit state handles (#870) +- SEP-2260 server requests must associate with a client request (#873) +- SEP-2549 follow-up: client-side TTL-honoring cache (#974) --- -## Informational (not scored for tiering) +## Completed -These draft/extension scenarios are tracked but do not count toward tier advancement: +### Tier 1 requirements -| Scenario | Tag | Status | -|---|---|---| -| `auth/resource-mismatch` | draft | ❌ Failed | -| `auth/client-credentials-jwt` | extension | ❌ Failed — JWT `aud` claim verification error | -| `auth/client-credentials-basic` | extension | ✅ Passed | -| `auth/cross-app-access-complete-flow` | extension | ❌ Failed — sends `authorization_code` grant instead of `jwt-bearer` | +- [x] **v3.0.0 stable released** (2026-07-28) — MRTR, SEP-2549 cache hints, SEP-2243 + standard headers, SEP-2575 stateless MCP, and SEP-2106 relaxations +- [x] 2025-11-25 server conformance 100% (30/30) +- [x] 2025-11-25 client conformance 100% +- [x] 2026-07-28 server conformance 100% (30/30 dated) +- [x] 2026-07-28 client conformance 100% (dated) +- [x] Issue triage ≥90% within 2 business days (95.2%, 20/21) with the full SEP-1730 + label taxonomy (bug, enhancement, question, needs confirmation, needs repro, + ready for work, good first issue, help wanted, P0–P3) +- [x] All P0 bugs resolved within 7 days (0 open; #741 resolved in 3 days). #815 was + reclassified from `P0` to `T-security` — it was a CVE/advisory-coordination task + for an already-shipped fix (PR #764, released in v1.4.0), not a critical-bug fix +- [x] `VERSIONING.md` — semver scheme, breaking-change definition, and communication policy +- [x] `DEPENDENCY_POLICY.md` — published dependency update policy (with `.github/dependabot.yml`) +- [x] `SECURITY.md` and Dependabot configuration +- [x] All 48 non-experimental features documented with examples in the README + (closed the last 6 gaps on 2026-07-30: audio results, resource-template reading, + resource-argument completion, ping, and the legacy HTTP+SSE non-goal writeup) + +### Spec implementation + +- [x] SEP-2322 MRTR (server scenarios + `sep-2322-client-request-state`) +- [x] SEP-2575 Make MCP Stateless (`server-stateless`) +- [x] SEP-2164 resource not found +- [x] SEP-2549 cache hints (`caching`) +- [x] SEP-2243 HTTP standardization (`http-header-validation`, standard headers) +- [x] DNS rebinding protection + +--- + +## Nice-to-have (scorecard hygiene, not required for Tier 1) + +- [ ] Add a top-level `CHANGELOG.md` (release notes are currently managed by release-plz). +- [ ] Add a top-level `CONTRIBUTING.md` (contributor docs currently live at `docs/CONTRIBUTE.MD`). diff --git a/VERSIONING.md b/VERSIONING.md new file mode 100644 index 0000000..e3117e0 --- /dev/null +++ b/VERSIONING.md @@ -0,0 +1,66 @@ +# Versioning Policy + +`rmcp` and `rmcp-macros` follow [Semantic Versioning 2.0.0](https://semver.org/), +the standard versioning scheme for the Rust/Cargo ecosystem. Given a version +`MAJOR.MINOR.PATCH`: + +- **MAJOR** — incremented for breaking changes to the public API. +- **MINOR** — incremented for new, backwards-compatible functionality. +- **PATCH** — incremented for backwards-compatible bug fixes. + +The `rmcp` and `rmcp-macros` crates are versioned together and released with the +same version number. + +## What counts as a breaking change + +We treat a change as breaking if it would cause code that compiled against the +previous release to fail to compile, or to change behavior in a way callers could +not reasonably anticipate. This follows the Cargo SemVer guidelines +([The Cargo Book — SemVer Compatibility](https://doc.rust-lang.org/cargo/reference/semver.html)). +Examples include: + +- Removing or renaming a public item (function, type, trait, module, variant, field). +- Changing a public function signature, trait method, or trait bound. +- Adding a required method to a public trait, or a field to a struct that callers + construct directly. (Public structs and enums are marked `#[non_exhaustive]` where + practical so that additive changes remain non-breaking.) +- Raising the Minimum Supported Rust Version (MSRV) — see below. + +Additive changes — new functions, new types, new enum variants on +`#[non_exhaustive]` enums, new optional Cargo features — are **not** breaking and +ship in MINOR releases. + +### Cargo features + +The public API is feature-gated. SemVer compatibility is evaluated against the +crate's **default features** and all features except `local`. + +## Pre-1.0 crates in the workspace + +The core crates (`rmcp`, `rmcp-macros`) are `>= 1.0.0` and follow the rules above. +Any workspace crate still in `0.x` follows Cargo's pre-1.0 convention, where the +`MINOR` version acts as the breaking-change signal (`0.MINOR.PATCH`). + +## Minimum Supported Rust Version (MSRV) + +The MSRV is declared via `rust-version` in `Cargo.toml`. Raising the MSRV is +considered a breaking change and is only done in a MINOR or MAJOR release with a +note in the release notes. + +## How breaking changes are communicated + +- Release automation is handled by [release-plz](https://release-plz.dev/), which + derives version bumps from [Conventional Commits](https://www.conventionalcommits.org/) + (`feat!:` / `fix!:` and `BREAKING CHANGE:` footers trigger a MAJOR bump) and runs + `cargo-semver-checks` to catch unintended API breaks. +- Every release is published as a tagged [GitHub Release](https://github.com/modelcontextprotocol/rust-sdk/releases) + with generated notes. +- Significant migrations (for example the 2.x → 3.x upgrade) are accompanied by a + migration guide linked from the release notes and the README. + +## Yanking + +A published version that is later found to be broken or to contain a security issue +may be [yanked](https://doc.rust-lang.org/cargo/commands/cargo-yank.html) from +crates.io. Yanking prevents new dependents from selecting the version but does not +delete it; a fixed PATCH release is published alongside. diff --git a/clippy.toml b/clippy.toml index 9388e00..b7c01ac 100644 --- a/clippy.toml +++ b/clippy.toml @@ -1,3 +1,2 @@ -msrv = "1.85" too-many-arguments-threshold = 10 check-private-items = false diff --git a/conformance/Cargo.toml b/conformance/Cargo.toml index 42a5e85..e8de6f0 100644 --- a/conformance/Cargo.toml +++ b/conformance/Cargo.toml @@ -18,6 +18,9 @@ rmcp = { path = "../crates/rmcp", features = [ "client", "elicitation", "auth", + "auth-client-credentials-jwt", + "auth-enterprise-managed", + "request-state", "transport-streamable-http-server", "transport-streamable-http-client-reqwest", ] } @@ -29,8 +32,7 @@ tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } axum = { version = "0.8", features = ["macros"] } anyhow = "1" +oauth2 = { version = "5.0", default-features = false } reqwest = { version = "0.13", features = ["json"] } urlencoding = "2" url = "2" -p256 = { version = "0.13", features = ["ecdsa"] } -base64 = "0.22" diff --git a/conformance/expected-failures-extensions.yaml b/conformance/expected-failures-extensions.yaml new file mode 100644 index 0000000..0ebbeef --- /dev/null +++ b/conformance/expected-failures-extensions.yaml @@ -0,0 +1,26 @@ +# Known failures for informational extension scenarios in +# @modelcontextprotocol/conformance CONFORMANCE_VERSION. +# +# Extensions are not selected by a `--spec-version` run and do not count toward +# SDK tiering. CI runs them separately so a green versioned-spec suite does not +# hide extension gaps. +# +# This is a strict baseline: +# - an unlisted failure fails the build +# - a listed scenario that starts passing fails the build as a stale entry +# +# When bumping CONFORMANCE_VERSION, review the available extension and +# pending scenarios and update this file deliberately. + +# The SEP-2663 Tasks Extension server scenarios (tracked in #868) all pass and +# were removed from this baseline. `tasks-status-notifications` remains +# upstream-skipped pending the subscriptions/listen rewrite; CI will fail if it +# becomes active but does not pass. +server: [] + +client: + # Informational OAuth extension scenarios. + - auth/enterprise-managed-authorization + - auth/dpop + - auth/dpop-nonce + - auth/wif-jwt-bearer diff --git a/conformance/src/bin/client.rs b/conformance/src/bin/client.rs index 8dabff0..aef2c27 100644 --- a/conformance/src/bin/client.rs +++ b/conformance/src/bin/client.rs @@ -1,10 +1,21 @@ +#![expect( + deprecated, + reason = "The conformance suite still exercises deprecated sampling scenarios" +)] + +use anyhow::Context; +use oauth2::{ClientSecret, RefreshToken}; use rmcp::{ - ClientHandler, ErrorData, RoleClient, ServiceExt, + ClientHandler, ClientLifecycleMode, ClientServiceExt, ErrorData, RoleClient, ServiceExt, model::*, service::RequestContext, transport::{ AuthClient, AuthorizationManager, StreamableHttpClientTransport, - auth::{AuthorizationCallback, OAuthState}, + auth::{ + AuthorizationCallback, AuthorizationRequest, ClientCredentialsConfig, + InMemoryCredentialStore, JwtSigningAlgorithm, OAuthState, default_oauth_http_client, + enterprise::{EmaAuthorizationServer, EmaClientAuthentication, EmaExchangeRequest}, + }, streamable_http_client::StreamableHttpClientTransportConfig, }, }; @@ -13,8 +24,17 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; // ─── Context parsed from MCP_CONFORMANCE_CONTEXT ──────────────────────────── +#[derive(Debug, Default, serde::Deserialize)] +struct ConformanceToolCall { + name: String, + #[serde(default)] + arguments: Option>, +} + #[derive(Debug, Default, serde::Deserialize)] struct ConformanceContext { + #[serde(default, alias = "toolCalls")] + tool_calls: Vec, #[serde(default)] client_id: Option, #[serde(default)] @@ -24,6 +44,17 @@ struct ConformanceContext { private_key_pem: Option, #[serde(default)] signing_algorithm: Option, + // enterprise-managed-authorization-refresh-token + #[serde(default)] + idp_client_id: Option, + #[serde(default)] + idp_client_secret: Option, + #[serde(default)] + idp_refresh_token: Option, + #[serde(default)] + idp_issuer: Option, + #[serde(default)] + idp_token_endpoint: Option, } fn load_context() -> ConformanceContext { @@ -43,8 +74,8 @@ impl ClientHandler for BasicClientHandler {} struct ElicitationDefaultsClientHandler; impl ClientHandler for ElicitationDefaultsClientHandler { - fn get_info(&self) -> ClientInfo { - let mut info = ClientInfo::default(); + fn get_info(&self) -> ClientConfig { + let mut info = ClientConfig::default(); info.capabilities.elicitation = Some( ElicitationCapability::new() .with_form(FormElicitationCapability::new().with_schema_validation(true)), @@ -132,8 +163,8 @@ impl ClientHandler for ElicitationDefaultsClientHandler { struct FullClientHandler; impl ClientHandler for FullClientHandler { - fn get_info(&self) -> ClientInfo { - let mut info = ClientInfo::default(); + fn get_info(&self) -> ClientConfig { + let mut info = ClientConfig::default(); info.capabilities.elicitation = Some( ElicitationCapability::new() .with_form(FormElicitationCapability::new().with_schema_validation(true)), @@ -180,27 +211,76 @@ impl ClientHandler for FullClientHandler { const CIMD_CLIENT_METADATA_URL: &str = "https://conformance-test.local/client-metadata.json"; const REDIRECT_URI: &str = "http://localhost:3000/callback"; +const SCOPE_STEP_UP_INITIAL_SCOPES: &[&str] = &["mcp:basic"]; +const SCOPE_STEP_UP_ESCALATED_SCOPES: &[&str] = &["mcp:basic", "mcp:write"]; -/// Perform the headless OAuth authorization-code flow. +/// Attempt the real connection unauthenticated and return the server's +/// `WWW-Authenticate` challenge from the 401 — the reactive discovery +/// trigger. /// -/// 1. Discover metadata, register (or use CIMD), get auth URL -/// 2. Fetch the auth URL with redirect:manual → extract code from Location header -/// 3. Exchange code for token -/// 4. Return an `AuthClient` wrapping `reqwest::Client` +/// `None` (server accepted the unauthenticated connection, which is then +/// closed cleanly) is a legitimate outcome, not an error: the scope-step-up +/// and scope-retry-limit mocks allow unauthenticated `initialize` and only +/// enforce authorization on tool calls. +async fn initialize_challenge( + server_url: &str, + lifecycle: ClientLifecycleMode, +) -> anyhow::Result> { + let transport = StreamableHttpClientTransport::from_uri(server_url); + match BasicClientHandler + .serve_with_lifecycle(transport, lifecycle) + .await + { + Ok(client) => { + client.cancel().await.ok(); + Ok(None) + } + Err(error) => match error.auth_challenge() { + Some(challenge) => Ok(Some(challenge.to_string())), + None => Err(error.into()), + }, + } +} + +fn with_optional_challenge( + request: AuthorizationRequest, + challenge: Option, +) -> AuthorizationRequest { + match challenge { + Some(challenge) => request.with_challenge(challenge), + None => request, + } +} + +/// Perform the headless OAuth authorization-code flow, reactively: +/// +/// 1. Attempt the real connection; take the 401's WWW-Authenticate challenge +/// 2. Discover from the challenge, register (or use CIMD), get auth URL +/// 3. Fetch the auth URL with redirect:manual → extract code from Location header +/// 4. Exchange code for token +/// 5. Return an `AuthClient` wrapping `reqwest::Client` async fn perform_oauth_flow( server_url: &str, _ctx: &ConformanceContext, ) -> anyhow::Result> { + // Always the discover lifecycle here (not `conformance_lifecycle()`): + // this flow serves `run_auth_client`, whose 2026-07-28 auth mocks require + // the per-request MCP-Protocol-Version negotiation. + let challenge = initialize_challenge( + server_url, + ClientLifecycleMode::Discover { + preferred_versions: preferred_protocol_versions(), + }, + ) + .await?; let mut oauth = OAuthState::new(server_url, None).await?; // Discover + register + get auth URL + let request = AuthorizationRequest::new(REDIRECT_URI) + .with_client_name("conformance-client") + .with_client_metadata_url(CIMD_CLIENT_METADATA_URL); oauth - .start_authorization_with_metadata_url( - &[], - REDIRECT_URI, - Some("conformance-client"), - Some(CIMD_CLIENT_METADATA_URL), - ) + .start_authorization(with_optional_challenge(request, challenge)) .await?; let auth_url = oauth.get_authorization_url().await?; @@ -235,50 +315,38 @@ async fn perform_oauth_flow( Ok(AuthClient::new(reqwest::Client::default(), am)) } -/// Like `perform_oauth_flow` but uses pre-registered client credentials. +/// Like `perform_oauth_flow` but uses pre-registered client credentials, +/// exercising the SDK's high-level `OAuthState` path (no DCR). async fn perform_oauth_flow_preregistered( server_url: &str, client_id: &str, client_secret: &str, ) -> anyhow::Result> { - let mut manager = AuthorizationManager::new(server_url).await?; - let metadata = manager.discover_metadata().await?; - manager.set_metadata(metadata); + let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?; + let mut oauth = OAuthState::new(server_url, None).await?; - // Configure with pre-registered credentials - let config = rmcp::transport::auth::OAuthClientConfig::new(client_id, REDIRECT_URI) + let request = AuthorizationRequest::new(REDIRECT_URI) + .with_preregistered_client(client_id) .with_client_secret(client_secret); - manager.configure_client(config)?; + oauth + .start_authorization(with_optional_challenge(request, challenge)) + .await?; - let scopes = manager.select_scopes(None, &[]); - let scope_refs: Vec<&str> = scopes.iter().map(|s| s.as_str()).collect(); - let auth_url = manager.get_authorization_url(&scope_refs).await?; + let auth_url = oauth.get_authorization_url().await?; + let callback = headless_authorize(&auth_url).await?; + oauth + .handle_callback_with_issuer( + &callback.code, + &callback.csrf_token, + callback.issuer.as_deref(), + ) + .await?; - // Headless redirect - let http = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build()?; - let resp = http.get(&auth_url).send().await?; - let location = resp - .headers() - .get("location") - .and_then(|v| v.to_str().ok()) - .ok_or_else(|| anyhow::anyhow!("No Location header"))?; - let redirect_url = url::Url::parse(location)?; - let code = redirect_url - .query_pairs() - .find(|(k, _)| k == "code") - .map(|(_, v)| v.to_string()) - .ok_or_else(|| anyhow::anyhow!("No code"))?; - let state = redirect_url - .query_pairs() - .find(|(k, _)| k == "state") - .map(|(_, v)| v.to_string()) - .ok_or_else(|| anyhow::anyhow!("No state"))?; + let am = oauth + .into_authorization_manager() + .ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?; - manager.exchange_code_for_token(&code, &state).await?; - - Ok(AuthClient::new(reqwest::Client::default(), manager)) + Ok(AuthClient::new(reqwest::Client::default(), am)) } /// Run the standard auth flow, then connect and exercise the server. @@ -290,7 +358,17 @@ async fn run_auth_client(server_url: &str, ctx: &ConformanceContext) -> anyhow:: StreamableHttpClientTransportConfig::with_uri(server_url), ); - let client = BasicClientHandler.serve(transport).await?; + // The 2026-07-28 auth mocks require the modern per-request lifecycle + // (MCP-Protocol-Version header on every request), so negotiate via the + // discover lifecycle rather than the legacy initialize handshake. + let client = BasicClientHandler + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: preferred_protocol_versions(), + }, + ) + .await?; tracing::debug!("Connected (authenticated)"); let tools = client.list_tools(Default::default()).await?; @@ -314,15 +392,14 @@ async fn run_auth_scope_step_up_client( server_url: &str, _ctx: &ConformanceContext, ) -> anyhow::Result<()> { - // First auth + let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?; let mut oauth = OAuthState::new(server_url, None).await?; + let request = AuthorizationRequest::new(REDIRECT_URI) + .with_scopes(SCOPE_STEP_UP_INITIAL_SCOPES.iter().copied()) + .with_client_name("conformance-client") + .with_client_metadata_url(CIMD_CLIENT_METADATA_URL); oauth - .start_authorization_with_metadata_url( - &[], - REDIRECT_URI, - Some("conformance-client"), - Some(CIMD_CLIENT_METADATA_URL), - ) + .start_authorization(with_optional_challenge(request, challenge)) .await?; let auth_url = oauth.get_authorization_url().await?; @@ -345,7 +422,9 @@ async fn run_auth_scope_step_up_client( StreamableHttpClientTransportConfig::with_uri(server_url), ); - let client = BasicClientHandler.serve(transport).await?; + let client = BasicClientHandler + .serve_with_lifecycle(transport, conformance_lifecycle()) + .await?; let tools = client.list_tools(Default::default()).await?; tracing::debug!("Listed {} tools", tools.tools.len()); @@ -365,16 +444,13 @@ async fn run_auth_scope_step_up_client( // Drop old client, re-auth with upgraded scopes client.cancel().await.ok(); - // Re-do the full flow; the server will give us the right scopes - // on the second authorization request. let mut oauth2 = OAuthState::new(server_url, None).await?; - // Pass the escalated scope hint oauth2 - .start_authorization_with_metadata_url( - &[], - REDIRECT_URI, - Some("conformance-client"), - Some(CIMD_CLIENT_METADATA_URL), + .start_authorization( + AuthorizationRequest::new(REDIRECT_URI) + .with_scopes(SCOPE_STEP_UP_ESCALATED_SCOPES.iter().copied()) + .with_client_name("conformance-client") + .with_client_metadata_url(CIMD_CLIENT_METADATA_URL), ) .await?; let auth_url2 = oauth2.get_authorization_url().await?; @@ -387,16 +463,20 @@ async fn run_auth_scope_step_up_client( ) .await?; - let am2 = oauth2.into_authorization_manager().unwrap(); + let am2 = oauth2.into_authorization_manager().ok_or_else(|| { + anyhow::anyhow!("Missing authorization manager after step-up") + })?; let auth_client2 = AuthClient::new(reqwest::Client::default(), am2); let transport2 = StreamableHttpClientTransport::with_client( auth_client2, StreamableHttpClientTransportConfig::with_uri(server_url), ); - let client2 = BasicClientHandler.serve(transport2).await?; - let _ = client2 + let client2 = BasicClientHandler + .serve_with_lifecycle(transport2, conformance_lifecycle()) + .await?; + client2 .call_tool(call_tool_params(tool.name.clone(), args)) - .await; + .await?; client2.cancel().await.ok(); return Ok(()); } @@ -414,16 +494,15 @@ async fn run_auth_scope_retry_limit_client( ) -> anyhow::Result<()> { let max_retries = 3u32; let mut attempt = 0u32; + let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?; loop { let mut oauth = OAuthState::new(server_url, None).await?; + let request = AuthorizationRequest::new(REDIRECT_URI) + .with_client_name("conformance-client") + .with_client_metadata_url(CIMD_CLIENT_METADATA_URL); oauth - .start_authorization_with_metadata_url( - &[], - REDIRECT_URI, - Some("conformance-client"), - Some(CIMD_CLIENT_METADATA_URL), - ) + .start_authorization(with_optional_challenge(request, challenge.clone())) .await?; let auth_url = oauth.get_authorization_url().await?; let callback = headless_authorize(&auth_url).await?; @@ -435,7 +514,9 @@ async fn run_auth_scope_retry_limit_client( ) .await?; - let am = oauth.into_authorization_manager().unwrap(); + let am = oauth + .into_authorization_manager() + .ok_or_else(|| anyhow::anyhow!("Missing authorization manager"))?; let auth_client = AuthClient::new(reqwest::Client::default(), am); let transport = StreamableHttpClientTransport::with_client( auth_client, @@ -443,7 +524,18 @@ async fn run_auth_scope_retry_limit_client( ); let client = BasicClientHandler.serve(transport).await?; - let tools = client.list_tools(Default::default()).await?; + let tools = match client.list_tools(Default::default()).await { + Ok(tools) => tools, + Err(err) => { + tracing::info!( + "Scope retry limit scenario stopped after authorization attempt {}: {}", + attempt + 1, + err + ); + client.cancel().await.ok(); + return Ok(()); + } + }; let mut got_403 = false; for tool in &tools.tools { @@ -467,12 +559,75 @@ async fn run_auth_scope_retry_limit_client( attempt += 1; if attempt >= max_retries { tracing::info!("Reached retry limit ({max_retries}), giving up"); - return Err(anyhow::anyhow!("Scope retry limit reached")); + return Ok(()); } } Ok(()) } +async fn migration_token( + server_url: &str, + store: &InMemoryCredentialStore, +) -> anyhow::Result { + let mut manager = AuthorizationManager::new(server_url).await?; + manager.set_credential_store(store.clone()); + + if manager.initialize_from_store().await? { + return Ok(manager.get_access_token().await?); + } + + let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?; + let resolution = manager + .resolve_metadata_from_challenge(challenge.as_deref()) + .await?; + manager.set_metadata(resolution.metadata); + manager + .register_client("conformance-client", REDIRECT_URI, &[]) + .await?; + + let scopes = manager.select_scopes(None, &[]); + let scope_refs: Vec<&str> = scopes.iter().map(|s| s.as_str()).collect(); + let auth_url = manager.get_authorization_url(&scope_refs).await?; + let callback = headless_authorize(&auth_url).await?; + manager + .exchange_code_for_token_with_issuer( + &callback.code, + &callback.csrf_token, + callback.issuer.as_deref(), + ) + .await?; + + Ok(manager.get_access_token().await?) +} + +async fn run_auth_server_migration_client( + server_url: &str, + _ctx: &ConformanceContext, +) -> anyhow::Result<()> { + let store = InMemoryCredentialStore::new(); + let http = reqwest::Client::new(); + let body = json!({"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}}); + + let mut token = migration_token(server_url, &store).await?; + for _ in 0..3 { + let resp = http + .post(server_url) + .header( + "MCP-Protocol-Version", + conformance_protocol_version().as_str(), + ) + .bearer_auth(&token) + .json(&body) + .send() + .await?; + if resp.status() == reqwest::StatusCode::UNAUTHORIZED { + token = migration_token(server_url, &store).await?; + } + } + + Ok(()) +} + /// Auth flow with pre-registered credentials (from context). async fn run_auth_preregistered_client( server_url: &str, @@ -524,9 +679,12 @@ async fn run_client_credentials_basic( .unwrap_or("conformance-test-secret"); let mut manager = AuthorizationManager::new(server_url).await?; - let metadata = manager.discover_metadata().await?; - let token_endpoint = metadata.token_endpoint.clone(); - manager.set_metadata(metadata); + let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?; + let resolution = manager + .resolve_metadata_from_challenge(challenge.as_deref()) + .await?; + let token_endpoint = resolution.metadata.token_endpoint.clone(); + manager.set_metadata(resolution.metadata); let http = reqwest::Client::new(); let resp = http @@ -569,49 +727,43 @@ async fn run_client_credentials_jwt( ) -> anyhow::Result<()> { let client_id = ctx .client_id - .as_deref() - .unwrap_or("conformance-test-client"); - let _pem = ctx + .clone() + .unwrap_or_else(|| "conformance-test-client".to_string()); + let signing_key = ctx .private_key_pem - .as_deref() - .ok_or_else(|| anyhow::anyhow!("Missing private_key_pem"))?; - let _alg = ctx + .as_ref() + .ok_or_else(|| anyhow::anyhow!("Missing private_key_pem"))? + .as_bytes() + .to_vec(); + let signing_algorithm = match ctx .signing_algorithm .as_deref() - .ok_or_else(|| anyhow::anyhow!("Missing signing_algorithm"))?; - - // Discover metadata to get token endpoint - let mut manager = AuthorizationManager::new(server_url).await?; - let metadata = manager.discover_metadata().await?; - let token_endpoint = metadata.token_endpoint.clone(); - manager.set_metadata(metadata); - - // Build JWT assertion - // Parse the PEM private key - let key = openssl_free_ec_sign(_pem, client_id, &token_endpoint)?; - - let http = reqwest::Client::new(); - let form_body = format!( - "grant_type=client_credentials&client_assertion_type={}&client_assertion={}", - urlencoding::encode("urn:ietf:params:oauth:client-assertion-type:jwt-bearer"), - urlencoding::encode(&key), - ); - let resp = http - .post(&token_endpoint) - .header("content-type", "application/x-www-form-urlencoded") - .body(form_body) - .send() - .await?; - - let token_resp: serde_json::Value = resp.json().await?; - let access_token = token_resp["access_token"] - .as_str() - .ok_or_else(|| anyhow::anyhow!("No access_token: {}", token_resp))?; + .ok_or_else(|| anyhow::anyhow!("Missing signing_algorithm"))? + { + "RS256" => JwtSigningAlgorithm::RS256, + "RS384" => JwtSigningAlgorithm::RS384, + "RS512" => JwtSigningAlgorithm::RS512, + "ES256" => JwtSigningAlgorithm::ES256, + "ES384" => JwtSigningAlgorithm::ES384, + algorithm => anyhow::bail!("Unsupported signing_algorithm: {algorithm}"), + }; + let config = ClientCredentialsConfig::PrivateKeyJwt { + client_id, + signing_key, + signing_algorithm, + token_endpoint_audience: None, + scopes: vec![], + resource: Some(server_url.to_string()), + }; + let mut oauth_state = OAuthState::new(server_url, None).await?; + oauth_state.authenticate_client_credentials(config).await?; + let manager = oauth_state + .into_authorization_manager() + .ok_or_else(|| anyhow::anyhow!("Client credentials flow did not authorize"))?; let transport = StreamableHttpClientTransport::with_client( - reqwest::Client::default(), - StreamableHttpClientTransportConfig::with_uri(server_url) - .auth_header(access_token.to_string()), + AuthClient::new(reqwest::Client::default(), manager), + StreamableHttpClientTransportConfig::with_uri(server_url), ); let client = BasicClientHandler.serve(transport).await?; @@ -627,71 +779,64 @@ async fn run_client_credentials_jwt( Ok(()) } -/// Minimal ES256 JWT signing without heavy deps. -/// We use ring or pure-Rust approach. For simplicity, use the p256 + base64 crates -/// that are already transitive deps of oauth2. -fn openssl_free_ec_sign(pem: &str, client_id: &str, audience: &str) -> anyhow::Result { - use std::time::{SystemTime, UNIX_EPOCH}; +/// Exchange the fixture's IdP refresh token, then exercise authenticated MCP access. +async fn run_ema_refresh_token_client( + server_url: &str, + ctx: &ConformanceContext, +) -> anyhow::Result<()> { + let manager = AuthorizationManager::new(server_url).await?; + let metadata = manager.resolve_metadata().await?.metadata; + let idp = EmaAuthorizationServer::new( + ctx.idp_issuer.as_deref().context("Missing idp_issuer")?, + ctx.idp_token_endpoint + .as_deref() + .context("Missing idp_token_endpoint")?, + ctx.idp_client_id + .as_deref() + .context("Missing idp_client_id")?, + ) + .with_client_authentication(EmaClientAuthentication::ClientSecretBasic( + ClientSecret::new( + ctx.idp_client_secret + .clone() + .context("Missing idp_client_secret")?, + ), + )); + let resource_as = EmaAuthorizationServer::new( + metadata + .issuer + .context("Missing authorization server issuer")?, + metadata.token_endpoint, + ctx.client_id.as_deref().context("Missing client_id")?, + ) + .with_client_authentication(EmaClientAuthentication::ClientSecretBasic( + ClientSecret::new(ctx.client_secret.clone().context("Missing client_secret")?), + )); + let refresh_token = RefreshToken::new( + ctx.idp_refresh_token + .clone() + .context("Missing idp_refresh_token")?, + ); + let http = default_oauth_http_client()?; + let token = EmaExchangeRequest::new(idp, resource_as, server_url, &refresh_token) + .with_scopes(manager.select_scopes(None, &[])) + .exchange(&http, &http) + .await?; - // Decode PEM → DER - let pem_body = pem - .lines() - .filter(|l| !l.starts_with("-----")) - .collect::(); - let der = base64_decode(&pem_body)?; - - // Parse PKCS#8 DER to get the raw EC private key bytes - // PKCS#8 for EC P-256: the raw 32-byte key is at the end of the structure - let raw_key = extract_ec_private_key(&der)?; - - let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(); - let header = base64url_encode(br#"{"alg":"ES256","typ":"JWT"}"#); - let payload_json = serde_json::json!({ - "iss": client_id, - "sub": client_id, - "aud": audience, - "iat": now, - "exp": now + 300, - "jti": format!("jti-{}", now), - }); - let payload = base64url_encode(payload_json.to_string().as_bytes()); - let signing_input = format!("{}.{}", header, payload); - - // Sign with p256 - let secret_key = p256::ecdsa::SigningKey::from_bytes(raw_key.as_slice().into()) - .map_err(|e| anyhow::anyhow!("Invalid EC key: {}", e))?; - use p256::ecdsa::signature::Signer; - let sig: p256::ecdsa::Signature = secret_key.sign(signing_input.as_bytes()); - let sig_bytes = sig.to_bytes(); - let sig_b64 = base64url_encode(&sig_bytes); - - Ok(format!("{}.{}", signing_input, sig_b64)) -} - -fn base64url_encode(data: &[u8]) -> String { - use base64::Engine; - base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data) -} - -fn base64_decode(s: &str) -> anyhow::Result> { - use base64::Engine; - Ok(base64::engine::general_purpose::STANDARD.decode(s.trim())?) -} - -/// Extract the raw 32-byte EC private key from a PKCS#8 DER blob. -fn extract_ec_private_key(der: &[u8]) -> anyhow::Result> { - // PKCS#8 wraps an ECPrivateKey. We look for the octet string containing - // the 32-byte private key. A simple heuristic: find 0x04 0x20 (OCTET STRING, len 32) - // followed by exactly 32 bytes that form the key. - // More robust: parse ASN.1. But for conformance testing this suffices. - for i in 0..der.len().saturating_sub(33) { - if der[i] == 0x04 && der[i + 1] == 0x20 && i + 34 <= der.len() { - return Ok(der[i + 2..i + 34].to_vec()); - } + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(server_url) + .auth_header(token.access_token.secret()), + ); + let client = BasicClientHandler + .serve_with_lifecycle(transport, conformance_lifecycle()) + .await?; + let tools = client.list_tools(Default::default()).await?; + for tool in tools.tools { + let args = build_tool_arguments(&tool); + client.call_tool(call_tool_params(tool.name, args)).await?; } - Err(anyhow::anyhow!( - "Could not extract 32-byte EC private key from PKCS#8 DER" - )) + client.cancel().await?; + Ok(()) } /// Cross-app access flow (SEP-1046 extension). @@ -780,16 +925,53 @@ async fn run_basic_client(server_url: &str) -> anyhow::Result<()> { Ok(()) } -async fn run_tools_call_client(server_url: &str) -> anyhow::Result<()> { +async fn run_tools_call_client(server_url: &str, ctx: &ConformanceContext) -> anyhow::Result<()> { + run_tools_call_client_with_lifecycle(server_url, ctx, conformance_lifecycle()).await +} + +async fn run_discover_tools_call_client( + server_url: &str, + ctx: &ConformanceContext, +) -> anyhow::Result<()> { + run_tools_call_client_with_lifecycle( + server_url, + ctx, + ClientLifecycleMode::Discover { + preferred_versions: preferred_protocol_versions(), + }, + ) + .await +} + +async fn run_tools_call_client_with_lifecycle( + server_url: &str, + ctx: &ConformanceContext, + lifecycle: ClientLifecycleMode, +) -> anyhow::Result<()> { let transport = StreamableHttpClientTransport::from_uri(server_url); - let client = FullClientHandler.serve(transport).await?; + let client = FullClientHandler + .serve_with_lifecycle(transport, lifecycle) + .await?; let tools = client.list_tools(Default::default()).await?; - for tool in &tools.tools { - let args = build_tool_arguments(tool); - let _ = client - .call_tool(call_tool_params(tool.name.clone(), args)) - .await?; + + if ctx.tool_calls.is_empty() { + for tool in &tools.tools { + let args = build_tool_arguments(tool); + client + .call_tool(call_tool_params(tool.name.clone(), args)) + .await?; + } + } else { + for tool_call in &ctx.tool_calls { + client + .call_tool(call_tool_params( + tool_call.name.clone().into(), + tool_call.arguments.clone(), + )) + .await?; + } } + client.cancel().await?; Ok(()) } @@ -811,6 +993,58 @@ async fn run_elicitation_defaults_client(server_url: &str) -> anyhow::Result<()> Ok(()) } +fn conformance_protocol_version() -> ProtocolVersion { + std::env::var("MCP_CONFORMANCE_PROTOCOL_VERSION") + .ok() + .and_then(|version| serde_json::from_value(Value::String(version)).ok()) + .unwrap_or(ProtocolVersion::V_2025_11_25) +} + +fn conformance_lifecycle() -> ClientLifecycleMode { + if conformance_protocol_version().as_str() >= ProtocolVersion::V_2026_07_28.as_str() { + ClientLifecycleMode::Discover { + preferred_versions: preferred_protocol_versions(), + } + } else { + ClientLifecycleMode::Initialize + } +} + +/// Preferred protocol versions for discover-lifecycle negotiation: the +/// runner-provided version first, then all other known versions newest-first. +fn preferred_protocol_versions() -> Vec { + let mut preferred_versions = vec![conformance_protocol_version()]; + for version in ProtocolVersion::KNOWN_VERSIONS.iter().rev() { + if !preferred_versions.contains(version) { + preferred_versions.push(version.clone()); + } + } + preferred_versions +} + +/// Runs scenarios through the discover lifecycle and Streamable HTTP transport. +async fn run_discover_client(server_url: &str) -> anyhow::Result<()> { + let preferred_versions = preferred_protocol_versions(); + let transport = StreamableHttpClientTransport::from_uri(server_url); + let client = FullClientHandler + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { preferred_versions }, + ) + .await?; + + let tools = client.list_tools(Default::default()).await?; + tracing::debug!("Listed {} tools", tools.tools.len()); + for tool in &tools.tools { + let result = client + .call_tool(CallToolRequestParams::new(tool.name.clone())) + .await; + tracing::debug!("Called {}: {:?}", tool.name, result.is_ok()); + } + client.cancel().await?; + Ok(()) +} + async fn run_sse_retry_client(server_url: &str) -> anyhow::Result<()> { let transport = StreamableHttpClientTransport::from_uri(server_url); let client = BasicClientHandler.serve(transport).await?; @@ -848,14 +1082,50 @@ async fn main() -> anyhow::Result<()> { tracing::info!("Running scenario '{}' against {}", scenario, server_url); - match scenario.as_str() { + // Safety net: some harness servers intentionally misbehave (e.g. reply + // with an id-less error instead of answering a request), which would + // leave the client waiting forever. Exit on our own before the harness's + // 30s client timeout so it never has to kill us (which has been observed + // to wedge the harness process in CI). + let timeout_secs: u64 = std::env::var("MCP_CONFORMANCE_TIMEOUT_SECS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(25); + tokio::time::timeout( + std::time::Duration::from_secs(timeout_secs), + run_scenario(&scenario, &server_url, &ctx), + ) + .await + .map_err(|_| anyhow::anyhow!("Scenario '{scenario}' timed out after {timeout_secs}s"))??; + + Ok(()) +} + +async fn run_scenario( + scenario: &str, + server_url: &str, + ctx: &ConformanceContext, +) -> anyhow::Result<()> { + match scenario { // Non-auth scenarios - "initialize" => run_basic_client(&server_url).await?, - "tools_call" => run_tools_call_client(&server_url).await?, + "initialize" => run_basic_client(server_url).await?, + // SEP-2106: the scenario serves a tool whose schema carries a network + // `$ref`; the check passes when the client lists tools without + // dereferencing (fetching) that URL. A plain connect → list_tools → + // close is sufficient; the scenario's mock server does not implement + // the discover lifecycle, so `run_discover_client` hangs against it. + "json-schema-ref-no-deref" => run_basic_client(server_url).await?, + "tools_call" => run_tools_call_client(server_url, ctx).await?, "elicitation-sep1034-client-defaults" => { - run_elicitation_defaults_client(&server_url).await? + run_elicitation_defaults_client(server_url).await? + } + "sse-retry" => run_sse_retry_client(server_url).await?, + "request-metadata" | "sep-2322-client-request-state" => { + run_discover_client(server_url).await? + } + "http-standard-headers" | "http-custom-headers" | "http-invalid-tool-headers" => { + run_discover_tools_call_client(server_url, ctx).await? } - "sse-retry" => run_sse_retry_client(&server_url).await?, // Auth scenarios - standard OAuth flow "auth/metadata-default" @@ -870,21 +1140,42 @@ async fn main() -> anyhow::Result<()> { | "auth/token-endpoint-auth-post" | "auth/token-endpoint-auth-none" | "auth/2025-03-26-oauth-metadata-backcompat" - | "auth/2025-03-26-oauth-endpoint-fallback" => run_auth_client(&server_url, &ctx).await?, + | "auth/2025-03-26-oauth-endpoint-fallback" + // Offline access scope handling: positive/negative variants both run + // the well-behaved flow; the referee inspects the requested scopes. + | "auth/offline-access-scope" + | "auth/offline-access-not-supported" + // SEP-2468 (RFC 9207 iss / RFC 8414 §3.3 issuer-echo). The client + // captures `iss` from the authorization redirect and passes it to the + // callback handler; the SDK validates internally. Positive scenarios + // proceed to the token endpoint; negative scenarios error out (the + // referee sets `allowClientError`). + | "auth/iss-supported" + | "auth/iss-not-advertised" + | "auth/iss-supported-missing" + | "auth/iss-wrong-issuer" + | "auth/iss-unexpected" + | "auth/iss-normalized" + | "auth/metadata-issuer-mismatch" => run_auth_client(server_url, ctx).await?, // Auth - scope step-up - "auth/scope-step-up" => run_auth_scope_step_up_client(&server_url, &ctx).await?, + "auth/scope-step-up" => run_auth_scope_step_up_client(server_url, ctx).await?, // Auth - scope retry limit - "auth/scope-retry-limit" => run_auth_scope_retry_limit_client(&server_url, &ctx).await?, + "auth/scope-retry-limit" => run_auth_scope_retry_limit_client(server_url, ctx).await?, + + // Auth - authorization server migration (SEP-2352) + "auth/authorization-server-migration" => { + run_auth_server_migration_client(server_url, ctx).await? + } // Auth - pre-registration - "auth/pre-registration" => run_auth_preregistered_client(&server_url, &ctx).await?, + "auth/pre-registration" => run_auth_preregistered_client(server_url, ctx).await?, // Auth - resource mismatch (should fail to auth → pass) "auth/resource-mismatch" => { // Try to auth; it should fail because PRM resource doesn't match - match run_auth_client(&server_url, &ctx).await { + match run_auth_client(server_url, ctx).await { Ok(_) => { tracing::warn!("Auth succeeded despite resource mismatch!"); } @@ -895,25 +1186,48 @@ async fn main() -> anyhow::Result<()> { } // Auth - client credentials - "auth/client-credentials-basic" => run_client_credentials_basic(&server_url, &ctx).await?, - "auth/client-credentials-jwt" => run_client_credentials_jwt(&server_url, &ctx).await?, + "auth/client-credentials-basic" => run_client_credentials_basic(server_url, ctx).await?, + "auth/client-credentials-jwt" => run_client_credentials_jwt(server_url, ctx).await?, + + // Auth - enterprise-managed authorization with a refresh-token subject + "auth/enterprise-managed-authorization-refresh-token" => { + run_ema_refresh_token_client(server_url, ctx).await? + } // Auth - cross-app access "auth/cross-app-access-complete-flow" => { - run_cross_app_access_client(&server_url, &ctx).await? + run_cross_app_access_client(server_url, ctx).await? } - _ => { - tracing::warn!("Unknown scenario '{}', trying auth flow", scenario); - match run_auth_client(&server_url, &ctx).await { - Ok(_) => {} - Err(e) => { - tracing::debug!("Auth flow failed for unknown scenario: {e}"); - run_basic_client(&server_url).await? - } - } - } + unknown => anyhow::bail!("Unsupported conformance scenario: {unknown}"), } Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn conformance_context_accepts_camel_case_tool_calls() { + let context: ConformanceContext = serde_json::from_value(json!({ + "toolCalls": [{ + "name": "test_custom_headers", + "arguments": { "region": "us-west1" } + }] + })) + .expect("valid conformance context"); + + assert_eq!( + context.tool_calls.first().map(|tool_call| ( + tool_call.name.as_str(), + tool_call + .arguments + .as_ref() + .and_then(|arguments| arguments.get("region")), + )), + Some(("test_custom_headers", Some(&json!("us-west1")))) + ); + } +} diff --git a/conformance/src/bin/server.rs b/conformance/src/bin/server.rs index b0b0d63..5c4b5c8 100644 --- a/conformance/src/bin/server.rs +++ b/conformance/src/bin/server.rs @@ -1,10 +1,17 @@ #![allow(deprecated)] -use std::{collections::HashSet, sync::Arc}; +use std::{ + collections::{HashMap, HashSet}, + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, +}; use rmcp::{ ErrorData, RoleServer, ServerHandler, model::*, - service::RequestContext, + service::{RequestContext, SubscriptionContext, SubscriptionSink}, + task_manager::{TaskExit, TaskManager, TaskOptions}, transport::{ StreamableHttpServerConfig, StreamableHttpService, streamable_http_server::session::local::LocalSessionManager, @@ -18,6 +25,7 @@ use tracing_subscriber::EnvFilter; const TEST_IMAGE_DATA: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="; // Small base64-encoded WAV (silence) const TEST_AUDIO_DATA: &str = "UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA="; +const CACHE_TTL_MS: u64 = 60_000; /// Helper to convert a serde_json::Value (must be an object) into a JsonObject fn json_object(v: Value) -> JsonObject { @@ -27,43 +35,752 @@ fn json_object(v: Value) -> JsonObject { } } +fn custom_header_tool() -> Tool { + Tool::new( + "test_custom_header", + "Validates SEP-2243 custom parameter headers", + json_object(json!({ + "type": "object", + "properties": { + "value": { "type": "string", "x-mcp-header": "Value" } + }, + "required": ["value"] + })), + ) +} + +/// Signing key for SEP-2322 `requestState` sealing. A fixed key is fine for a +/// conformance harness; real servers must load a secret out of clients' reach. +const REQUEST_STATE_KEY: &[u8] = b"rust-sdk-conformance-request-state-key!!"; +const _: () = assert!( + REQUEST_STATE_KEY.len() >= RequestStateCodec::MIN_KEY_LENGTH, + "REQUEST_STATE_KEY is shorter than RequestStateCodec::MIN_KEY_LENGTH", +); + #[derive(Clone)] struct ConformanceServer { - subscriptions: Arc>>, + legacy_resource_subscriptions: Arc>>, + subscriptions: Arc>>, + next_subscription: Arc, log_level: Arc>, + request_state_codec: RequestStateCodec, + tasks: TaskManager, } impl ConformanceServer { fn new() -> Self { Self { - subscriptions: Arc::new(Mutex::new(HashSet::new())), + legacy_resource_subscriptions: Arc::new(Mutex::new(HashSet::new())), + subscriptions: Arc::new(Mutex::new(HashMap::new())), + next_subscription: Arc::new(AtomicU64::new(0)), log_level: Arc::new(Mutex::new(LoggingLevel::Debug)), + request_state_codec: RequestStateCodec::new_unchecked(REQUEST_STATE_KEY), + tasks: TaskManager::new(), + } + } +} + +// ─── SEP-2663 Tasks extension fixtures ────────────────────────────────────── + +/// Fixture tools required by the Tasks extension conformance scenarios. +const TASK_FIXTURE_TOOLS: &[&str] = &[ + "greet", + "slow_compute", + "failing_job", + "protocol_error_job", + "confirm_delete", + "multi_input", + "test_tool_with_task", +]; + +/// Tools that are registered as task-supporting. `greet` is deliberately +/// sync-only. +const TASK_SUPPORTING_TOOLS: &[&str] = &[ + "slow_compute", + "failing_job", + "protocol_error_job", + "confirm_delete", + "multi_input", + "test_tool_with_task", +]; + +/// Tools that cannot be serviced without returning a `CreateTaskResult`: +/// calling them from a client that did not declare the tasks extension is +/// rejected with -32021 before the tool body runs (SEP-2663 §Required +/// Capabilities). `failing_job` and `test_tool_with_task` are registered +/// this way for the required-task-error and MRTR-composition scenarios; +/// `confirm_delete` and `multi_input` must park on in-task elicitation, so +/// they have no synchronous fallback either. +const TASK_REQUIRED_TOOLS: &[&str] = &[ + "failing_job", + "test_tool_with_task", + "confirm_delete", + "multi_input", +]; + +fn task_fixture_tool(name: &str) -> Tool { + let (description, schema) = match name { + "greet" => ( + "Sync-only greeting fixture (SEP-2663)", + json!({ + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"] + }), + ), + "slow_compute" => ( + "Task-supporting fixture: sleeps `seconds` then returns a result (SEP-2663)", + json!({ + "type": "object", + "properties": { + "seconds": { "type": "number" }, + "label": { "type": "string" } + } + }), + ), + "failing_job" => ( + "Task-supporting fixture (task support: required): returns a tool execution error (SEP-2663)", + json!({ "type": "object", "properties": {} }), + ), + "protocol_error_job" => ( + "Task-supporting fixture: fails with a protocol-level error (SEP-2663)", + json!({ "type": "object", "properties": {} }), + ), + "confirm_delete" => ( + "Task-supporting fixture: parks on a single elicitation inputRequest (SEP-2663)", + json!({ + "type": "object", + "properties": { "filename": { "type": "string" } } + }), + ), + "multi_input" => ( + "Task-supporting fixture: parks on two parallel elicitation inputRequests (SEP-2663)", + json!({ "type": "object", "properties": {} }), + ), + "test_tool_with_task" => ( + "MRTR round 1 gathers user_name, round 2 escalates to a task (SEP-2663 composition)", + json!({ "type": "object", "properties": {} }), + ), + other => panic!("unknown task fixture tool: {other}"), + }; + Tool::new(name.to_string(), description, json_object(schema)) +} + +// ─── SEP-2322 MRTR (InputRequiredResult) helpers ──────────────────────────── + +fn mrtr_elicitation_request(message: &str, properties: Value, required: Value) -> InputRequest { + InputRequest::Elicitation(ElicitRequest::new( + ElicitRequestParams::FormElicitationParams { + meta: None, + message: message.into(), + requested_schema: serde_json::from_value(json!({ + "type": "object", + "properties": properties, + "required": required, + })) + .expect("valid elicitation schema"), + }, + )) +} + +fn mrtr_sampling_request(prompt: &str) -> InputRequest { + InputRequest::CreateMessage(CreateMessageRequest::new(CreateMessageRequestParams::new( + vec![SamplingMessage::user_text(prompt)], + 100, + ))) +} + +fn mrtr_list_roots_request() -> InputRequest { + InputRequest::ListRoots(ListRootsRequest::default()) +} + +/// An input response is usable when it is a JSON object (an `ElicitResult`, +/// `CreateMessageResult`, or `ListRootsResult` shape). Anything else (e.g. a +/// bare number) is treated as missing so the server re-requests it. +fn mrtr_response<'a>( + responses: Option<&'a InputResponses>, + key: &str, +) -> Option<&'a serde_json::Map> { + responses + .and_then(|r| r.get(key)) + .and_then(Value::as_object) +} + +impl ConformanceServer { + fn mrtr_tampered_state_error() -> ErrorData { + ErrorData::invalid_params("requestState failed integrity verification", None) + } + + /// SEP-2663 task fixture tools. The server decides per request whether to + /// materialize a task: task-supporting tools create one when the client + /// declared the tasks extension capability; otherwise they fall through to + /// synchronous execution (except task-*required* tools, which reject with + /// -32021). + async fn call_task_fixture_tool( + &self, + request: CallToolRequestParams, + cx: &RequestContext, + ) -> Result { + let client_supports_tasks = cx + .client_capabilities() + .is_some_and(|caps| caps.supports_tasks()); + let name = request.name.as_ref(); + let args = request.arguments.clone().unwrap_or_default(); + + if TASK_REQUIRED_TOOLS.contains(&name) && !client_supports_tasks { + // SEP-2663 §Required Capabilities: this tool cannot be serviced + // without returning CreateTaskResult. + return Err(ErrorData::missing_required_client_capability( + ClientCapabilities::builder().enable_tasks().build(), + )); + } + + let create_task = client_supports_tasks && TASK_SUPPORTING_TOOLS.contains(&name); + + match name { + "greet" => { + let who = args.get("name").and_then(Value::as_str).unwrap_or("friend"); + Ok( + CallToolResult::success(vec![ContentBlock::text(format!("Hello, {who}!"))]) + .into(), + ) + } + + "slow_compute" => { + let seconds = args.get("seconds").and_then(Value::as_f64).unwrap_or(1.0); + let label = args + .get("label") + .and_then(Value::as_str) + .unwrap_or("compute") + .to_string(); + if create_task { + // The lifecycle scenario requires slow_compute to settle + // to `cancelled` when tasks/cancel arrives while running; + // cancellation is cooperative, so honor it explicitly. + let task = self.tasks.spawn(TaskOptions::default(), move |ctx| { + Box::pin(async move { + tokio::select! { + _ = ctx.cancelled() => Err(TaskExit::Cancelled), + _ = tokio::time::sleep( + std::time::Duration::from_secs_f64(seconds), + ) => Ok(CallToolResult::success(vec![ContentBlock::text( + format!("slow_compute({label}) done after {seconds}s"), + )])), + } + }) + }); + Ok(CreateTaskResult::new(task).into()) + } else { + tokio::time::sleep(std::time::Duration::from_secs_f64(seconds)).await; + Ok(CallToolResult::success(vec![ContentBlock::text(format!( + "slow_compute({label}) done after {seconds}s" + ))]) + .into()) + } + } + + "failing_job" => { + // Tool execution error: surfaces as status "completed" with + // result.isError = true when run as a task. + let work = || async { + tokio::time::sleep(std::time::Duration::from_millis(1000)).await; + Ok(CallToolResult::error(vec![ContentBlock::text( + "failing_job: intentional tool execution error", + )])) + }; + if create_task { + let task = self.tasks.spawn(TaskOptions::default(), move |_ctx| { + Box::pin(async move { work().await.map_err(TaskExit::Error) }) + }); + Ok(CreateTaskResult::new(task).into()) + } else { + Ok(work().await?.into()) + } + } + + "protocol_error_job" => { + // Protocol-level failure: surfaces as status "failed" with an + // inlined `error` object when run as a task. + let work = || async { + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + Err(ErrorData::internal_error( + "protocol_error_job: intentional protocol-level failure", + None, + )) + }; + if create_task { + let task = self.tasks.spawn(TaskOptions::default(), move |_ctx| { + Box::pin(async move { work().await.map_err(TaskExit::Error) }) + }); + Ok(CreateTaskResult::new(task).into()) + } else { + work().await.map(CallToolResponse::from) + } + } + + "confirm_delete" => { + let filename = args + .get("filename") + .and_then(Value::as_str) + .unwrap_or("file.txt") + .to_string(); + let task = self.tasks.spawn(TaskOptions::default(), move |ctx| { + Box::pin(async move { + let response = ctx + .request_input( + "confirm", + mrtr_elicitation_request( + &format!("Delete {filename}?"), + json!({ "confirm": { "type": "boolean" } }), + json!(["confirm"]), + ), + ) + .await?; + let confirmed = response + .get("content") + .and_then(|c| c.get("confirm")) + .and_then(Value::as_bool) + .unwrap_or(false); + Ok(CallToolResult::success(vec![ContentBlock::text(format!( + "confirm_delete({filename}): confirmed = {confirmed}" + ))])) + }) + }); + Ok(CreateTaskResult::new(task).into()) + } + + "multi_input" => { + let task = self.tasks.spawn(TaskOptions::default(), move |ctx| { + Box::pin(async move { + // Fan out two elicitation requests in parallel so two + // keys are pending at once (partial fulfillment check). + let first = ctx.request_input( + "input-a", + mrtr_elicitation_request( + "Provide value A", + json!({ "value": { "type": "string" } }), + json!(["value"]), + ), + ); + let second = ctx.request_input( + "input-b", + mrtr_elicitation_request( + "Provide value B", + json!({ "value": { "type": "string" } }), + json!(["value"]), + ), + ); + let (a, b) = tokio::join!(first, second); + let (a, b) = (a?, b?); + let get = |v: &Value| { + v.get("content") + .and_then(|c| c.get("value")) + .and_then(Value::as_str) + .unwrap_or("(none)") + .to_string() + }; + Ok(CallToolResult::success(vec![ContentBlock::text(format!( + "multi_input: a = {}, b = {}", + get(&a), + get(&b) + ))])) + }) + }); + Ok(CreateTaskResult::new(task).into()) + } + + "test_tool_with_task" => { + // SEP-2663 MRTR → Tasks composition. Round 1 (no inputResponses) + // is a plain MRTR InputRequiredResult; round 2 escalates to a + // task whose result reflects the gathered user_name. + match mrtr_response(request.input_responses.as_ref(), "user_name") { + None => { + let mut requests = InputRequests::new(); + requests.insert( + "user_name".into(), + mrtr_elicitation_request( + "What is your name?", + json!({ "name": { "type": "string" } }), + json!(["name"]), + ), + ); + Ok(InputRequiredResult::from_input_requests(requests).into()) + } + Some(response) => { + let user_name = response + .get("content") + .and_then(|c| c.get("name")) + .and_then(Value::as_str) + .unwrap_or("friend") + .to_string(); + let task = self.tasks.spawn(TaskOptions::default(), move |_ctx| { + Box::pin(async move { + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + Ok(CallToolResult::success(vec![ContentBlock::text(format!( + "Hello, {user_name}! (async)" + ))])) + }) + }); + Ok(CreateTaskResult::new(task).into()) + } + } + } + + other => Err(ErrorData::invalid_params( + format!("Unknown task fixture tool: {other}"), + None, + )), + } + } + + /// SEP-2322 test tools. Each returns an `InputRequiredResult` until the + /// client retries with the expected `inputResponses` (and, where used, the + /// echoed `requestState`). + /// + /// `meta` is the request's `_meta`, which the service loop moves out of the + /// params and into the `RequestContext`. + async fn call_mrtr_tool( + &self, + request: CallToolRequestParams, + meta: &RequestMetaObject, + ) -> Result { + let responses = request.input_responses.as_ref(); + match request.name.as_ref() { + "test_input_required_result_elicitation" => { + match mrtr_response(responses, "user_name") { + Some(response) => { + let name = response + .get("content") + .and_then(|c| c.get("name")) + .and_then(Value::as_str) + .unwrap_or("friend"); + Ok(CallToolResult::success(vec![ContentBlock::text(format!( + "Hello, {name}!" + ))]) + .into()) + } + // Initial call, or a retry with missing/invalid responses: + // (re-)request the input per the SEP's recommendation. + None => { + let mut requests = InputRequests::new(); + requests.insert( + "user_name".into(), + mrtr_elicitation_request( + "What is your name?", + json!({ "name": { "type": "string" } }), + json!(["name"]), + ), + ); + Ok(InputRequiredResult::from_input_requests(requests).into()) + } + } + } + + "test_input_required_result_sampling" => { + match mrtr_response(responses, "capital_question") { + Some(response) => { + let text = response + .get("content") + .and_then(|c| c.get("text")) + .and_then(Value::as_str) + .unwrap_or("(no sampling text)"); + Ok(CallToolResult::success(vec![ContentBlock::text(format!( + "Sampling response: {text}" + ))]) + .into()) + } + None => { + let mut requests = InputRequests::new(); + requests.insert( + "capital_question".into(), + mrtr_sampling_request("What is the capital of France?"), + ); + Ok(InputRequiredResult::from_input_requests(requests).into()) + } + } + } + + "test_input_required_result_list_roots" => { + match mrtr_response(responses, "client_roots") { + Some(response) => { + let roots = response + .get("roots") + .and_then(Value::as_array) + .map(|roots| { + roots + .iter() + .filter_map(|r| r.get("uri").and_then(Value::as_str)) + .collect::>() + .join(", ") + }) + .unwrap_or_default(); + Ok(CallToolResult::success(vec![ContentBlock::text(format!( + "Client roots: [{roots}]" + ))]) + .into()) + } + None => { + let mut requests = InputRequests::new(); + requests.insert("client_roots".into(), mrtr_list_roots_request()); + Ok(InputRequiredResult::from_input_requests(requests).into()) + } + } + } + + "test_input_required_result_request_state" + | "test_input_required_result_tampered_state" => { + match request.request_state.as_deref() { + // Initial call: request confirmation and seal our progress. + None => { + let sealed = self + .request_state_codec + .seal_json(&json!({ "stage": "confirm" })) + .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; + let mut requests = InputRequests::new(); + requests.insert( + "confirm".into(), + mrtr_elicitation_request( + "Please confirm", + json!({ "ok": { "type": "boolean" } }), + json!(["ok"]), + ), + ); + Ok(InputRequiredResult::new(Some(requests), Some(sealed)).into()) + } + // Retry: the echoed state is untrusted input and MUST pass + // integrity verification before we act on it. + Some(sealed) => { + self.request_state_codec + .open(sealed) + .map_err(|_| Self::mrtr_tampered_state_error())?; + Ok( + CallToolResult::success(vec![ContentBlock::text( + "Confirmed: state-ok", + )]) + .into(), + ) + } + } + } + + "test_input_required_result_multiple_inputs" => { + if let Some(sealed) = request.request_state.as_deref() { + self.request_state_codec + .open(sealed) + .map_err(|_| Self::mrtr_tampered_state_error())?; + } + let all_present = mrtr_response(responses, "user_name").is_some() + && mrtr_response(responses, "greeting").is_some() + && mrtr_response(responses, "client_roots").is_some(); + if all_present && request.request_state.is_some() { + Ok( + CallToolResult::success(vec![ContentBlock::text("All inputs received")]) + .into(), + ) + } else { + let sealed = self + .request_state_codec + .seal_json(&json!({ "stage": "gather" })) + .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; + let mut requests = InputRequests::new(); + requests.insert( + "user_name".into(), + mrtr_elicitation_request( + "What is your name?", + json!({ "name": { "type": "string" } }), + json!(["name"]), + ), + ); + requests.insert( + "greeting".into(), + mrtr_sampling_request("Generate a greeting"), + ); + requests.insert("client_roots".into(), mrtr_list_roots_request()); + Ok(InputRequiredResult::new(Some(requests), Some(sealed)).into()) + } + } + + "test_input_required_result_multi_round" => { + let round = match request.request_state.as_deref() { + None => 0, + Some(sealed) => { + let state: Value = self + .request_state_codec + .open_json(sealed) + .map_err(|_| Self::mrtr_tampered_state_error())?; + state.get("round").and_then(Value::as_i64).unwrap_or(0) + } + }; + match round { + 0 => { + let sealed = self + .request_state_codec + .seal_json(&json!({ "round": 1 })) + .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; + let mut requests = InputRequests::new(); + requests.insert( + "step1".into(), + mrtr_elicitation_request( + "Step 1: What is your name?", + json!({ "name": { "type": "string" } }), + json!(["name"]), + ), + ); + Ok(InputRequiredResult::new(Some(requests), Some(sealed)).into()) + } + 1 => { + let sealed = self + .request_state_codec + .seal_json(&json!({ "round": 2 })) + .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; + let mut requests = InputRequests::new(); + requests.insert( + "step2".into(), + mrtr_elicitation_request( + "Step 2: What is your favorite color?", + json!({ "color": { "type": "string" } }), + json!(["color"]), + ), + ); + Ok(InputRequiredResult::new(Some(requests), Some(sealed)).into()) + } + _ => Ok(CallToolResult::success(vec![ContentBlock::text( + "Multi-round flow complete", + )]) + .into()), + } + } + + "test_input_required_result_capabilities" => { + if responses.is_some() { + return Ok(CallToolResult::success(vec![ContentBlock::text( + "Capability-aware flow complete", + )]) + .into()); + } + // Per SEP-2322, only request inputs the client declared support + // for in `_meta['io.modelcontextprotocol/clientCapabilities']`. + let capabilities = meta.client_capabilities().unwrap_or_default(); + let mut requests = InputRequests::new(); + if capabilities.elicitation.is_some() { + requests.insert( + "user_name".into(), + mrtr_elicitation_request( + "What is your name?", + json!({ "name": { "type": "string" } }), + json!(["name"]), + ), + ); + } + if capabilities.sampling.is_some() { + requests.insert( + "greeting".into(), + mrtr_sampling_request("Generate a greeting"), + ); + } + if capabilities.roots.is_some() { + requests.insert("client_roots".into(), mrtr_list_roots_request()); + } + if requests.is_empty() { + Ok(CallToolResult::success(vec![ContentBlock::text( + "Client declared no MRTR-capable capabilities", + )]) + .into()) + } else { + Ok(InputRequiredResult::from_input_requests(requests).into()) + } + } + + _ => Err(ErrorData::invalid_params( + format!("Unknown tool: {}", request.name), + None, + )), } } } impl ServerHandler for ConformanceServer { - async fn initialize( - &self, - _request: InitializeRequestParams, - _cx: RequestContext, - ) -> Result { - Ok(InitializeResult::new( + fn get_tool(&self, name: &str) -> Option { + (name == "test_custom_header").then(custom_header_tool) + } + + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_prompts() + .enable_prompts_list_changed() .enable_resources() + .enable_resources_subscribe() + .enable_resources_list_changed() .enable_tools() + .enable_tool_list_changed() .enable_logging() + .enable_tasks() .build(), ) .with_server_info(Implementation::new("rust-conformance-server", "0.1.0")) - .with_instructions("Rust MCP conformance test server")) + .with_instructions("Rust MCP conformance test server") + } + + async fn initialize( + &self, + request: InitializeRequestParams, + _cx: RequestContext, + ) -> Result { + let info = self.get_info(); + Ok(InitializeResult::new(info.capabilities) + .with_protocol_version(request.protocol_version) + .with_server_info(info.server_info) + .with_instructions(info.instructions.unwrap_or_default())) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.supported_by(&self.get_info().capabilities)) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), ErrorData> { + let key = self.next_subscription.fetch_add(1, Ordering::Relaxed); + self.subscriptions + .lock() + .await + .insert(key, context.sink().clone()); + context.cancelled().await; + self.subscriptions.lock().await.remove(&key); + Ok(()) } async fn ping(&self, _cx: RequestContext) -> Result<(), ErrorData> { Ok(()) } + async fn get_task( + &self, + request: GetTaskParams, + _cx: RequestContext, + ) -> Result { + Ok(GetTaskResult::new(self.tasks.get_task(&request.task_id)?)) + } + + async fn update_task( + &self, + request: UpdateTaskParams, + _cx: RequestContext, + ) -> Result<(), ErrorData> { + self.tasks + .update_task(&request.task_id, request.input_responses) + } + + async fn cancel_task( + &self, + request: CancelTaskParams, + _cx: RequestContext, + ) -> Result<(), ErrorData> { + self.tasks.cancel_task(&request.task_id) + } + async fn list_tools( &self, _request: Option, @@ -180,6 +897,7 @@ impl ServerHandler for ConformanceServer { "type": "object", "$defs": { "address": { + "$anchor": "address", "type": "object", "properties": { "street": { "type": "string" }, @@ -191,6 +909,19 @@ impl ServerHandler for ConformanceServer { "name": { "type": "string" }, "address": { "$ref": "#/$defs/address" } }, + "allOf": [{ + "anyOf": [ + { "required": ["name"] }, + { "required": ["address"] } + ] + }], + "if": { "required": ["address"] }, + "then": { + "properties": { + "address": { "required": ["street"] } + } + }, + "else": { "required": ["name"] }, "additionalProperties": false })), ), @@ -202,21 +933,119 @@ impl ServerHandler for ConformanceServer { "properties": {} })), ), + custom_header_tool(), + Tool::new( + "test_trigger_tool_change", + "Triggers a tools/list_changed notification on matching subscriptions", + json_object(json!({ + "type": "object", + "properties": {} + })), + ), + Tool::new( + "test_trigger_prompt_change", + "Triggers a prompts/list_changed notification on matching subscriptions", + json_object(json!({ + "type": "object", + "properties": {} + })), + ), + Tool::new( + "test_missing_capability", + "Requires the sampling client capability", + json_object(json!({ + "type": "object", + "properties": {} + })), + ), + Tool::new( + "test_streaming_elicitation", + "Returns an input_required result containing an elicitation request", + json_object(json!({ + "type": "object", + "properties": {} + })), + ), + Tool::new( + "test_logging_tool", + "Emits notifications/message only when logLevel is requested", + json_object(json!({ + "type": "object", + "properties": {} + })), + ), ]; + // SEP-2322 MRTR test tools; all take no arguments. + let mrtr_tools = [ + ( + "test_input_required_result_elicitation", + "Requires an elicitation input via InputRequiredResult (SEP-2322)", + ), + ( + "test_input_required_result_sampling", + "Requires a sampling input via InputRequiredResult (SEP-2322)", + ), + ( + "test_input_required_result_list_roots", + "Requires a roots/list input via InputRequiredResult (SEP-2322)", + ), + ( + "test_input_required_result_request_state", + "Round-trips integrity-protected requestState (SEP-2322)", + ), + ( + "test_input_required_result_multiple_inputs", + "Requires elicitation + sampling + roots inputs in one round (SEP-2322)", + ), + ( + "test_input_required_result_multi_round", + "Drives multiple input_required rounds with evolving requestState (SEP-2322)", + ), + ( + "test_input_required_result_tampered_state", + "Rejects tampered requestState with a JSON-RPC error (SEP-2322)", + ), + ( + "test_input_required_result_capabilities", + "Only requests inputs for declared client capabilities (SEP-2322)", + ), + ]; + let tools = tools + .into_iter() + .chain(mrtr_tools.into_iter().map(|(name, description)| { + Tool::new( + name, + description, + json_object(json!({ "type": "object", "properties": {} })), + ) + })) + .chain( + TASK_FIXTURE_TOOLS + .iter() + .map(|name| task_fixture_tool(name)), + ) + .collect(); Ok(ListToolsResult { - meta: None, tools, - next_cursor: None, - }) + ..Default::default() + } + .with_ttl_ms(CACHE_TTL_MS) + .with_cache_scope(CacheScope::Public)) } async fn call_tool( &self, request: CallToolRequestParams, cx: RequestContext, - ) -> Result { + ) -> Result { + if request.name.starts_with("test_input_required_result_") { + return self.call_mrtr_tool(request, &cx.meta).await; + } + if TASK_FIXTURE_TOOLS.contains(&request.name.as_ref()) { + return self.call_task_fixture_tool(request, &cx).await; + } let args = request.arguments.unwrap_or_default(); - match request.name.as_ref() { + let result = match request.name.as_ref() { "test_simple_text" => Ok(CallToolResult::success(vec![ContentBlock::text( "This is a simple text response for testing.", )])), @@ -300,6 +1129,14 @@ impl ServerHandler for ConformanceServer { )])) } + "test_custom_header" => { + let value = args + .get("value") + .and_then(Value::as_str) + .ok_or_else(|| ErrorData::invalid_params("value must be a string", None))?; + Ok(CallToolResult::success(vec![ContentBlock::text(value)])) + } + "test_sampling" => { let prompt = args .get("prompt") @@ -527,11 +1364,87 @@ impl ServerHandler for ConformanceServer { )])) } + "test_trigger_tool_change" => { + let subscriptions = self + .subscriptions + .lock() + .await + .values() + .cloned() + .collect::>(); + for subscription in subscriptions { + let _ = subscription.notify_tool_list_changed().await; + } + Ok(CallToolResult::success(vec![ContentBlock::text( + "Tool list change triggered", + )])) + } + + "test_trigger_prompt_change" => { + let subscriptions = self + .subscriptions + .lock() + .await + .values() + .cloned() + .collect::>(); + for subscription in subscriptions { + let _ = subscription.notify_prompt_list_changed().await; + } + Ok(CallToolResult::success(vec![ContentBlock::text( + "Prompt list change triggered", + )])) + } + + "test_missing_capability" => { + let capabilities = cx.meta.client_capabilities().unwrap_or_default(); + if capabilities.sampling.is_none() { + return Err(ErrorData::missing_required_client_capability( + ClientCapabilities::builder().enable_sampling().build(), + )); + } + Ok(CallToolResult::success(vec![ContentBlock::text( + "Required capability declared", + )])) + } + + "test_streaming_elicitation" => { + let mut requests = InputRequests::new(); + requests.insert( + "streaming_elicitation".into(), + mrtr_elicitation_request( + "Provide a value", + json!({ "value": { "type": "string" } }), + json!(["value"]), + ), + ); + return Ok(InputRequiredResult::from_input_requests(requests).into()); + } + + "test_logging_tool" => { + if let Some(level) = cx.meta.log_level() { + let _ = cx + .peer + .notify_logging_message( + LoggingMessageNotificationParam::new( + level, + json!("logLevel was requested"), + ) + .with_logger("conformance-server"), + ) + .await; + } + Ok(CallToolResult::success(vec![ContentBlock::text( + "Logging tool completed", + )])) + } + _ => Err(ErrorData::invalid_params( format!("Unknown tool: {}", request.name), None, )), - } + }; + result.map(Into::into) } async fn list_resources( @@ -540,7 +1453,6 @@ impl ServerHandler for ConformanceServer { _cx: RequestContext, ) -> Result { Ok(ListResourcesResult { - meta: None, resources: vec![ Resource::new("test://static-text", "Static Text Resource") .with_description("A static text resource for testing") @@ -549,17 +1461,19 @@ impl ServerHandler for ConformanceServer { .with_description("A static binary/blob resource for testing") .with_mime_type("image/png"), ], - next_cursor: None, - }) + ..Default::default() + } + .with_ttl_ms(CACHE_TTL_MS) + .with_cache_scope(CacheScope::Public)) } async fn read_resource( &self, request: ReadResourceRequestParams, _cx: RequestContext, - ) -> Result { + ) -> Result { let uri = request.uri.as_str(); - match uri { + let result = match uri { "test://static-text" => Ok(ReadResourceResult::new(vec![ ResourceContents::TextResourceContents { uri: uri.into(), @@ -600,7 +1514,14 @@ impl ServerHandler for ConformanceServer { )) } } - } + }; + result + .map(|result| { + result + .with_ttl_ms(CACHE_TTL_MS) + .with_cache_scope(CacheScope::Public) + }) + .map(Into::into) } async fn list_resource_templates( @@ -609,14 +1530,15 @@ impl ServerHandler for ConformanceServer { _cx: RequestContext, ) -> Result { Ok(ListResourceTemplatesResult { - meta: None, resource_templates: vec![ ResourceTemplate::new("test://template/{id}/data", "Dynamic Resource") .with_description("A dynamic resource with parameter substitution") .with_mime_type("application/json"), ], - next_cursor: None, - }) + ..Default::default() + } + .with_ttl_ms(CACHE_TTL_MS) + .with_cache_scope(CacheScope::Public)) } async fn subscribe( @@ -624,7 +1546,7 @@ impl ServerHandler for ConformanceServer { request: SubscribeRequestParams, _cx: RequestContext, ) -> Result<(), ErrorData> { - let mut subs = self.subscriptions.lock().await; + let mut subs = self.legacy_resource_subscriptions.lock().await; subs.insert(request.uri.to_string()); Ok(()) } @@ -634,7 +1556,7 @@ impl ServerHandler for ConformanceServer { request: UnsubscribeRequestParams, _cx: RequestContext, ) -> Result<(), ErrorData> { - let mut subs = self.subscriptions.lock().await; + let mut subs = self.legacy_resource_subscriptions.lock().await; subs.remove(request.uri.as_str()); Ok(()) } @@ -645,7 +1567,6 @@ impl ServerHandler for ConformanceServer { _cx: RequestContext, ) -> Result { Ok(ListPromptsResult { - meta: None, prompts: vec![ Prompt::new( "test_simple_prompt", @@ -656,11 +1577,11 @@ impl ServerHandler for ConformanceServer { "test_prompt_with_arguments", Some("A test prompt that accepts arguments"), Some(vec![ - PromptArgument::new("name") - .with_description("The name to greet") + PromptArgument::new("arg1") + .with_description("First test argument") .with_required(true), - PromptArgument::new("style") - .with_description("The greeting style") + PromptArgument::new("arg2") + .with_description("Second test argument") .with_required(false), ]), ), @@ -674,17 +1595,56 @@ impl ServerHandler for ConformanceServer { Some("A test prompt that includes an image"), None, ), + Prompt::new( + "test_input_required_result_prompt", + Some( + "A prompt that requires elicitation input via InputRequiredResult (SEP-2322)", + ), + None, + ), ], - next_cursor: None, - }) + ..Default::default() + } + .with_ttl_ms(CACHE_TTL_MS) + .with_cache_scope(CacheScope::Public)) } async fn get_prompt( &self, request: GetPromptRequestParams, _cx: RequestContext, - ) -> Result { - match request.name.as_str() { + ) -> Result { + // SEP-2322: InputRequiredResult on a non-tool request (prompts/get). + if request.name == "test_input_required_result_prompt" { + return match mrtr_response(request.input_responses.as_ref(), "user_context") { + Some(response) => { + let context = response + .get("content") + .and_then(|c| c.get("context")) + .and_then(Value::as_str) + .unwrap_or("(no context)"); + Ok(GetPromptResult::new(vec![PromptMessage::new_text( + Role::User, + format!("Prompt with elicited context: {context}"), + )]) + .with_description("A prompt built from elicited context") + .into()) + } + None => { + let mut requests = InputRequests::new(); + requests.insert( + "user_context".into(), + mrtr_elicitation_request( + "What context should the prompt use?", + json!({ "context": { "type": "string" } }), + json!(["context"]), + ), + ); + Ok(InputRequiredResult::from_input_requests(requests).into()) + } + }; + } + let result = match request.name.as_str() { "test_simple_prompt" => Ok(GetPromptResult::new(vec![PromptMessage::new_text( Role::User, "This is a simple test prompt.", @@ -692,14 +1652,11 @@ impl ServerHandler for ConformanceServer { .with_description("A simple test prompt")), "test_prompt_with_arguments" => { let args = request.arguments.unwrap_or_default(); - let name = args.get("name").and_then(|v| v.as_str()).unwrap_or("World"); - let style = args - .get("style") - .and_then(|v| v.as_str()) - .unwrap_or("friendly"); + let arg1 = args.get("arg1").and_then(|v| v.as_str()).unwrap_or(""); + let arg2 = args.get("arg2").and_then(|v| v.as_str()).unwrap_or(""); Ok(GetPromptResult::new(vec![PromptMessage::new_text( Role::User, - format!("Please greet {} in a {} style.", name, style), + format!("Prompt with arguments: arg1='{}', arg2='{}'", arg1, arg2), )]) .with_description("A prompt with arguments")) } @@ -728,7 +1685,8 @@ impl ServerHandler for ConformanceServer { format!("Unknown prompt: {}", request.name), None, )), - } + }; + result.map(Into::into) } async fn complete( @@ -786,7 +1744,9 @@ async fn main() -> anyhow::Result<()> { tracing::info!("Starting conformance server on {}", bind_addr); let server = ConformanceServer::new(); - let config = StreamableHttpServerConfig::default(); + let config = StreamableHttpServerConfig::default() + .with_allowed_origins([format!("http://{bind_addr}")]) + .enforce_origin_validation(); let service = StreamableHttpService::new( move || Ok(server.clone()), LocalSessionManager::default().into(), @@ -801,3 +1761,26 @@ async fn main() -> anyhow::Result<()> { Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn server_exposes_custom_header_tool_for_transport_validation() { + let tool = ConformanceServer::new() + .get_tool("test_custom_header") + .expect("custom-header conformance tool"); + let value = Value::Object((*tool.input_schema).clone()); + + assert_eq!( + value.pointer("/properties/value/type"), + Some(&json!("string")) + ); + assert_eq!( + value.pointer("/properties/value/x-mcp-header"), + Some(&json!("Value")) + ); + assert_eq!(value.pointer("/required/0"), Some(&json!("value"))); + } +} diff --git a/crates/rmcp-macros/CHANGELOG.md b/crates/rmcp-macros/CHANGELOG.md index 242e591..6e35cf3 100644 --- a/crates/rmcp-macros/CHANGELOG.md +++ b/crates/rmcp-macros/CHANGELOG.md @@ -7,6 +7,99 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [3.4.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v3.3.0...rmcp-macros-v3.4.0) - 2026-09-15 + +### Added + +- *(model)* add ServerConfig and ClientConfig ([#1266](https://github.com/modelcontextprotocol/rust-sdk/pull/1266)) + +### Fixed + +- *(model)* deprecate ServerInfo and ClientInfo aliases ([#1156](https://github.com/modelcontextprotocol/rust-sdk/pull/1156)) + +## [3.3.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v3.2.0...rmcp-macros-v3.3.0) - 2026-09-10 + +### Added + +- *(macros)* reject empty tool_router ([#1233](https://github.com/modelcontextprotocol/rust-sdk/pull/1233)) + +## [3.2.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v3.1.4...rmcp-macros-v3.2.0) - 2026-08-31 + +### Added + +- add request-state key rotation ([#1128](https://github.com/modelcontextprotocol/rust-sdk/pull/1128)) + +### Fixed + +- allow concurrent streamable http requests ([#1186](https://github.com/modelcontextprotocol/rust-sdk/pull/1186)) + +## [3.1.3](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v3.1.2...rmcp-macros-v3.1.3) - 2026-08-17 + +### Fixed + +- time out auto discovery probe ([#1149](https://github.com/modelcontextprotocol/rust-sdk/pull/1149)) + +## [3.1.1](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v3.1.0...rmcp-macros-v3.1.1) - 2026-08-05 + +### Fixed + +- emit cache hints from handler macros ([#1120](https://github.com/modelcontextprotocol/rust-sdk/pull/1120)) + +### Other + +- upgrade darling and syn ([#1138](https://github.com/modelcontextprotocol/rust-sdk/pull/1138)) + +## [3.1.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v3.0.1...rmcp-macros-v3.1.0) - 2026-07-31 + +### Added + +- add strict stateless protocol metadata validation ([#1091](https://github.com/modelcontextprotocol/rust-sdk/pull/1091)) + +### Other + +- document the ping utility with examples ([#1106](https://github.com/modelcontextprotocol/rust-sdk/pull/1106)) +- complete Tier 1 feature docs and finalize roadmap ([#1101](https://github.com/modelcontextprotocol/rust-sdk/pull/1101)) +- *(conformance)* meeting requirements for tier 1 ([#1087](https://github.com/modelcontextprotocol/rust-sdk/pull/1087)) + +## [3.0.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v3.0.0-beta.5...rmcp-macros-v3.0.0) - 2026-07-28 + +### Other + +- release stable 3.0.0 + +## [3.0.0-beta.5](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v3.0.0-beta.4...rmcp-macros-v3.0.0-beta.5) - 2026-07-28 + +### Other + +- prepare for stable 3.0 release ([#1073](https://github.com/modelcontextprotocol/rust-sdk/pull/1073)) + +## [3.0.0-beta.2](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v3.0.0-beta.1...rmcp-macros-v3.0.0-beta.2) - 2026-07-24 + +### Fixed + +- [**breaking**] omit resultType for legacy protocol sessions ([#1038](https://github.com/modelcontextprotocol/rust-sdk/pull/1038)) + +### Other + +- declare and check MSRV ([#1034](https://github.com/modelcontextprotocol/rust-sdk/pull/1034)) + +## [3.0.0-beta.1](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v2.2.0...rmcp-macros-v3.0.0-beta.1) - 2026-07-23 + +### Added + +- add client-side TTL-honoring response cache (SEP-2549) ([#1025](https://github.com/modelcontextprotocol/rust-sdk/pull/1025)) +- [**breaking**] Implement SEP-2663 Tasks Extension ([#1020](https://github.com/modelcontextprotocol/rust-sdk/pull/1020)) +- add subscription listen streams (SEP-2575) ([#1000](https://github.com/modelcontextprotocol/rust-sdk/pull/1000)) +- add modern client lifecycle modes (SEP-2575) ([#995](https://github.com/modelcontextprotocol/rust-sdk/pull/995)) +- [**breaking**] implement SEP-2549 cache hints ([#889](https://github.com/modelcontextprotocol/rust-sdk/pull/889)) +- [**breaking**] add MRTR behavior support (SEP-2322) ([#929](https://github.com/modelcontextprotocol/rust-sdk/pull/929)) +- relax outputSchema to accept non-object JSON Schema types (SEP-2106) ([#895](https://github.com/modelcontextprotocol/rust-sdk/pull/895)) +- [**breaking**] add MRTR model types (SEP-2322) ([#915](https://github.com/modelcontextprotocol/rust-sdk/pull/915)) + +### Other + +- update for 2026-07-28 version ([#1032](https://github.com/modelcontextprotocol/rust-sdk/pull/1032)) + ## [2.0.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-macros-v1.8.0...rmcp-macros-v2.0.0) - 2026-06-27 ### Added diff --git a/crates/rmcp-macros/Cargo.toml b/crates/rmcp-macros/Cargo.toml index 6f645bd..9f0b3d3 100644 --- a/crates/rmcp-macros/Cargo.toml +++ b/crates/rmcp-macros/Cargo.toml @@ -5,6 +5,7 @@ name = "rmcp-macros" license = { workspace = true } version = { workspace = true } edition = { workspace = true } +rust-version = { workspace = true } repository = { workspace = true } homepage = { workspace = true } readme = { workspace = true } @@ -15,11 +16,11 @@ documentation = "https://docs.rs/rmcp-macros" proc-macro = true [dependencies] -syn = {version = "2", features = ["full"]} +syn = {version = "3", features = ["full"]} quote = "1" proc-macro2 = "1" serde_json = "1.0" -darling = { version = "0.23" } +darling = { version = "0.24" } [features] local = [] diff --git a/crates/rmcp-macros/README.md b/crates/rmcp-macros/README.md index cd9262e..adc8388 100644 --- a/crates/rmcp-macros/README.md +++ b/crates/rmcp-macros/README.md @@ -25,7 +25,6 @@ For **getting started** and **full MCP feature documentation**, see the [main RE | [`#[prompt]`][prompt] | Mark a function as an MCP prompt handler | | [`#[prompt_router]`][prompt_router] | Generate a prompt router from an impl block | | [`#[prompt_handler]`][prompt_handler] | Generate `get_prompt` and `list_prompts` handler methods | -| [`#[task_handler]`][task_handler] | Wire up the task lifecycle on top of an `OperationProcessor` | [tool]: https://docs.rs/rmcp-macros/latest/rmcp_macros/attr.tool.html [tool_router]: https://docs.rs/rmcp-macros/latest/rmcp_macros/attr.tool_router.html @@ -33,7 +32,6 @@ For **getting started** and **full MCP feature documentation**, see the [main RE [prompt]: https://docs.rs/rmcp-macros/latest/rmcp_macros/attr.prompt.html [prompt_router]: https://docs.rs/rmcp-macros/latest/rmcp_macros/attr.prompt_router.html [prompt_handler]: https://docs.rs/rmcp-macros/latest/rmcp_macros/attr.prompt_handler.html -[task_handler]: https://docs.rs/rmcp-macros/latest/rmcp_macros/attr.task_handler.html ## Quick Example diff --git a/crates/rmcp-macros/src/common.rs b/crates/rmcp-macros/src/common.rs index 0ca2fda..dc7e5bf 100644 --- a/crates/rmcp-macros/src/common.rs +++ b/crates/rmcp-macros/src/common.rs @@ -55,17 +55,15 @@ pub fn extract_doc_line( /// Returns the full Parameters type if found pub fn find_parameters_type_in_sig(sig: &Signature) -> Option> { sig.inputs.iter().find_map(|input| { - if let FnArg::Typed(pat_type) = input { - if let Type::Path(type_path) = &*pat_type.ty { - if type_path - .path - .segments - .last() - .is_some_and(|type_name| type_name.ident == "Parameters") - { - return Some(pat_type.ty.clone()); - } - } + if let FnArg::Typed(pat_type) = input + && let Type::Path(type_path) = &*pat_type.ty + && type_path + .path + .segments + .last() + .is_some_and(|type_name| type_name.ident == "Parameters") + { + return Some(pat_type.ty.clone()); } None }) diff --git a/crates/rmcp-macros/src/lib.rs b/crates/rmcp-macros/src/lib.rs index e721338..30433a8 100644 --- a/crates/rmcp-macros/src/lib.rs +++ b/crates/rmcp-macros/src/lib.rs @@ -7,7 +7,6 @@ mod common; mod prompt; mod prompt_handler; mod prompt_router; -mod task_handler; mod tool; mod tool_handler; mod tool_router; @@ -57,6 +56,7 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> TokenStream { /// | `router` | `Ident` | The name of the router function to be generated. Defaults to `tool_router`. | /// | `vis` | `Visibility` | The visibility of the generated router function. Defaults to empty. | /// | `server_handler` | `flag` | When set, also emits `#[::rmcp::tool_handler]` on `impl ServerHandler for Self` so you can omit a separate `#[tool_handler]` block. | +/// | `allow_empty` | `flag` | When set, accepts an impl block with no `#[tool]` fn. Without it, an empty router is a compile error. | /// /// ## Example /// @@ -100,7 +100,7 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> TokenStream { /// impl MyToolHandler { /// #[tool] /// fn my_tool_a() { -/// +/// /// } /// } /// } @@ -110,7 +110,7 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> TokenStream { /// impl MyToolHandler { /// #[tool] /// fn my_tool_b() { -/// +/// /// } /// } /// } @@ -123,6 +123,37 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> TokenStream { /// } /// } /// ``` +/// +/// ### Empty routers +/// +/// Collecting tools is this attribute's whole purpose, so an impl block with no `#[tool]` fn is a +/// compile error rather than a router that silently serves nothing. Pass `allow_empty` when that +/// is what you want: +/// +/// ```rust,ignore +/// #[tool_router(allow_empty)] +/// impl MyToolHandler {} +/// ``` +/// +/// The usual way to hit this by accident is a `macro_rules!` helper *inside* the impl block. An +/// attribute macro receives the unexpanded item, so `#[tool]` fns produced by such a helper are +/// invisible to `#[tool_router]`. Let the `macro_rules!` emit the whole annotated impl instead: +/// +/// ```rust,ignore +/// macro_rules! define_tools { +/// ($($name:ident => $description:literal),* $(,)?) => { +/// #[tool_router] +/// impl MyToolHandler { +/// $( +/// #[tool(description = $description)] +/// async fn $name(&self) -> String { stringify!($name).to_owned() } +/// )* +/// } +/// }; +/// } +/// +/// define_tools!(my_tool => "what my tool does"); +/// ``` #[proc_macro_attribute] pub fn tool_router(attr: TokenStream, input: TokenStream) -> TokenStream { tool_router::tool_router(attr.into(), input.into()) @@ -186,8 +217,8 @@ pub fn tool_router(attr: TokenStream, input: TokenStream) -> TokenStream { /// ```rust,ignore /// #[tool_handler] /// impl ServerHandler for MyToolHandler { -/// fn get_info(&self) -> ServerInfo { -/// ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) +/// fn get_info(&self) -> ServerConfig { +/// ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) /// } /// } /// ``` @@ -299,17 +330,3 @@ pub fn prompt_handler(attr: TokenStream, input: TokenStream) -> TokenStream { .unwrap_or_else(|err| err.to_compile_error()) .into() } - -/// # task_handler -/// -/// Generates basic task-handling methods (`enqueue_task` and `list_tasks`) for a server handler -/// using a shared \[`OperationProcessor`\]. The default processor expression assumes a -/// `self.processor` field holding an `Arc>`, but it can be customized -/// via `#[task_handler(processor = ...)]`. Because the macro captures `self` inside spawned -/// futures, the handler type must implement [`Clone`]. -#[proc_macro_attribute] -pub fn task_handler(attr: TokenStream, input: TokenStream) -> TokenStream { - task_handler::task_handler(attr.into(), input.into()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() -} diff --git a/crates/rmcp-macros/src/prompt.rs b/crates/rmcp-macros/src/prompt.rs index 3bf02d2..ea35bc1 100644 --- a/crates/rmcp-macros/src/prompt.rs +++ b/crates/rmcp-macros/src/prompt.rs @@ -132,13 +132,13 @@ pub fn prompt(attr: TokenStream, input: TokenStream) -> syn::Result // 3. make body: { Box::pin(async move { #body }) } let new_output = syn::parse2::({ let mut lt = quote! { 'static }; - if let Some(receiver) = fn_item.sig.receiver() { - if let Some((_, receiver_lt)) = receiver.reference.as_ref() { - if let Some(receiver_lt) = receiver_lt { - lt = quote! { #receiver_lt }; - } else { - lt = quote! { '_ }; - } + if let Some(receiver) = fn_item.sig.receiver() + && let syn::ReceiverKind::Reference(_, receiver_lt, _) = &receiver.kind + { + if let Some(receiver_lt) = receiver_lt { + lt = quote! { #receiver_lt }; + } else { + lt = quote! { '_ }; } } match &fn_item.sig.output { @@ -259,6 +259,20 @@ mod test { Ok(()) } + #[test] + fn test_async_prompt_preserves_receiver_lifetime() -> syn::Result<()> { + let attr = quote! {}; + let input = quote! { + async fn test_prompt_with_lifetime<'a>(&'a self) -> String { + "ok".to_string() + } + }; + let result = prompt(attr, input)?; + + assert!(result.to_string().contains("+ 'a")); + Ok(()) + } + #[test] fn test_async_prompt_local_omits_send() -> syn::Result<()> { let attr = quote! { local }; diff --git a/crates/rmcp-macros/src/prompt_handler.rs b/crates/rmcp-macros/src/prompt_handler.rs index 19ba438..6d957ec 100644 --- a/crates/rmcp-macros/src/prompt_handler.rs +++ b/crates/rmcp-macros/src/prompt_handler.rs @@ -35,7 +35,7 @@ pub fn prompt_handler(attr: TokenStream, input: TokenStream) -> syn::Result, - ) -> Result { + ) -> Result { let prompt_context = rmcp::handler::server::prompt::PromptContext::new( self, request.name, @@ -57,13 +57,20 @@ pub fn prompt_handler(attr: TokenStream, input: TokenStream) -> syn::Result, - _context: rmcp::service::RequestContext, + context: rmcp::service::RequestContext, ) -> Result { let prompts = #router_expr.list_all(); + let supports_cache_hints = context.protocol_version().is_some_and(|version| { + version >= rmcp::model::ProtocolVersion::V_2026_07_28 + }); Ok(rmcp::model::ListPromptsResult { + result_type: Some(rmcp::model::ResultType::COMPLETE), prompts, meta: #meta, next_cursor: None, + ttl_ms: supports_cache_hints.then_some(0), + cache_scope: supports_cache_hints + .then_some(rmcp::model::CacheScope::Public), }) } }; diff --git a/crates/rmcp-macros/src/task_handler.rs b/crates/rmcp-macros/src/task_handler.rs deleted file mode 100644 index 5815fd3..0000000 --- a/crates/rmcp-macros/src/task_handler.rs +++ /dev/null @@ -1,277 +0,0 @@ -use darling::{FromMeta, ast::NestedMeta}; -use proc_macro2::TokenStream; -use quote::{ToTokens, quote}; -use syn::{Expr, ImplItem, ItemImpl}; - -use crate::common::{has_method, has_sibling_handler}; - -#[derive(FromMeta)] -#[darling(default)] -struct TaskHandlerAttribute { - processor: Expr, -} - -impl Default for TaskHandlerAttribute { - fn default() -> Self { - Self { - processor: syn::parse2(quote! { self.processor }).expect("default processor expr"), - } - } -} - -pub fn task_handler(attr: TokenStream, input: TokenStream) -> syn::Result { - let attr_args = NestedMeta::parse_meta_list(attr)?; - let TaskHandlerAttribute { processor } = TaskHandlerAttribute::from_list(&attr_args)?; - let mut item_impl = syn::parse2::(input)?; - - if !has_method("list_tasks", &item_impl) { - let list_fn = quote! { - async fn list_tasks( - &self, - _request: Option, - _: rmcp::service::RequestContext, - ) -> Result { - let running_ids = (#processor).lock().await.list_running(); - let total = running_ids.len() as u64; - let tasks = running_ids - .into_iter() - .map(|task_id| { - let timestamp = rmcp::task_manager::current_timestamp(); - rmcp::model::Task::new( - task_id, - rmcp::model::TaskStatus::Working, - timestamp.clone(), - timestamp, - ) - }) - .collect::>(); - - Ok(rmcp::model::ListTasksResult::new(tasks)) - } - }; - item_impl.items.push(syn::parse2::(list_fn)?); - } - - if !has_method("enqueue_task", &item_impl) { - let enqueue_fn = quote! { - async fn enqueue_task( - &self, - request: rmcp::model::CallToolRequestParams, - context: rmcp::service::RequestContext, - ) -> Result { - use rmcp::task_manager::{ - current_timestamp, OperationDescriptor, OperationMessage, OperationResultTransport, - ToolCallTaskResult, - }; - let task_id = context.id.to_string(); - let operation_name = request.name.to_string(); - let future_request = request.clone(); - let future_context = context.clone(); - let server = self.clone(); - - let descriptor = OperationDescriptor::new(task_id.clone(), operation_name) - .with_context(context) - .with_client_request(rmcp::model::ClientRequest::CallToolRequest( - rmcp::model::Request::new(request), - )); - - let task_result_id = task_id.clone(); - let future = Box::pin(async move { - let result = server.call_tool(future_request, future_context).await; - Ok( - Box::new(ToolCallTaskResult::new(task_result_id, result)) - as Box, - ) - }); - - (#processor) - .lock() - .await - .submit_operation(OperationMessage::new(descriptor, future)) - .map_err(|err| rmcp::ErrorData::internal_error( - format!("failed to enqueue task: {err}"), - None, - ))?; - - let timestamp = current_timestamp(); - let task = rmcp::model::Task::new( - task_id, - rmcp::model::TaskStatus::Working, - timestamp.clone(), - timestamp, - ).with_status_message("Task accepted"); - - Ok(rmcp::model::CreateTaskResult::new(task)) - } - }; - item_impl.items.push(syn::parse2::(enqueue_fn)?); - } - - if !has_method("get_task_info", &item_impl) { - let get_info_fn = quote! { - async fn get_task_info( - &self, - request: rmcp::model::GetTaskParams, - _context: rmcp::service::RequestContext, - ) -> Result { - use rmcp::task_manager::current_timestamp; - let task_id = request.task_id.clone(); - let mut processor = (#processor).lock().await; - - // Check completed results first - let completed = processor.peek_completed().iter().rev().find(|r| r.descriptor.operation_id == task_id); - if let Some(completed_result) = completed { - // Determine Finished vs Failed - let status = match &completed_result.result { - Ok(boxed) => { - if let Some(tool) = boxed.as_any().downcast_ref::() { - match &tool.result { - Ok(_) => rmcp::model::TaskStatus::Completed, - Err(_) => rmcp::model::TaskStatus::Failed, - } - } else { - rmcp::model::TaskStatus::Completed - } - } - Err(_) => rmcp::model::TaskStatus::Failed, - }; - let timestamp = current_timestamp(); - let mut task = rmcp::model::Task::new( - task_id, - status, - timestamp.clone(), - timestamp, - ); - if let Some(ttl) = completed_result.descriptor.ttl { - task = task.with_ttl(ttl); - } - return Ok(rmcp::model::GetTaskResult::new(task)); - } - - // If not completed, check running - let running = processor.list_running(); - if running.into_iter().any(|id| id == task_id) { - let timestamp = current_timestamp(); - let task = rmcp::model::Task::new( - task_id, - rmcp::model::TaskStatus::Working, - timestamp.clone(), - timestamp, - ); - return Ok(rmcp::model::GetTaskResult::new(task)); - } - - Err(McpError::resource_not_found(format!("task not found: {}", task_id), None)) - } - }; - item_impl.items.push(syn::parse2::(get_info_fn)?); - } - - if !has_method("get_task_result", &item_impl) { - let get_result_fn = quote! { - async fn get_task_result( - &self, - request: rmcp::model::GetTaskPayloadParams, - _context: rmcp::service::RequestContext, - ) -> Result { - use std::time::Duration; - let task_id = request.task_id.clone(); - - loop { - // Scope the lock so we can await outside if needed - { - let mut processor = (#processor).lock().await; - - if let Some(task_result) = processor.take_completed_result(&task_id) { - match task_result.result { - Ok(boxed) => { - if let Some(tool) = boxed.as_any().downcast_ref::() { - match &tool.result { - Ok(call_tool) => { - let value = ::rmcp::serde_json::to_value(call_tool).unwrap_or_default(); - return Ok(rmcp::model::GetTaskPayloadResult::new(value)); - } - Err(err) => return Err(McpError::internal_error( - format!("task failed: {}", err), - None, - )), - } - } else { - return Err(McpError::internal_error("unsupported task result transport", None)); - } - } - Err(err) => return Err(McpError::internal_error( - format!("task execution error: {}", err), - None, - )), - } - } - - // Not completed yet: if not running, return not found - let running = processor.list_running(); - if !running.iter().any(|id| id == &task_id) { - return Err(McpError::resource_not_found(format!("task not found: {}", task_id), None)); - } - } - - tokio::time::sleep(Duration::from_millis(100)).await; - } - } - }; - item_impl - .items - .push(syn::parse2::(get_result_fn)?); - } - - if !has_method("cancel_task", &item_impl) { - let cancel_fn = quote! { - async fn cancel_task( - &self, - request: rmcp::model::CancelTaskParams, - _context: rmcp::service::RequestContext, - ) -> Result { - use rmcp::task_manager::current_timestamp; - let task_id = request.task_id; - let mut processor = (#processor).lock().await; - - if processor.cancel_task(&task_id) { - let timestamp = current_timestamp(); - let task = rmcp::model::Task::new( - task_id, - rmcp::model::TaskStatus::Cancelled, - timestamp.clone(), - timestamp, - ); - return Ok(rmcp::model::CancelTaskResult::new(task)); - } - - // If already completed, signal it's not cancellable - let exists_completed = processor.peek_completed().iter().any(|r| r.descriptor.operation_id == task_id); - if exists_completed { - return Err(McpError::invalid_request(format!("task already completed: {}", task_id), None)); - } - - Err(McpError::resource_not_found(format!("task not found: {}", task_id), None)) - } - }; - item_impl.items.push(syn::parse2::(cancel_fn)?); - } - - // Auto-generate get_info() if not already provided and no sibling tool/prompt handler - // will generate it (they take priority since they run as outer attributes). - if !has_method("get_info", &item_impl) - && !has_sibling_handler(&item_impl, "tool_handler") - && !has_sibling_handler(&item_impl, "prompt_handler") - { - let get_info_fn = crate::tool_handler::build_get_info( - &item_impl, - None, - None, - None, - crate::tool_handler::CallerCapability::Tasks, - )?; - item_impl.items.push(get_info_fn); - } - - Ok(item_impl.into_token_stream()) -} diff --git a/crates/rmcp-macros/src/tool.rs b/crates/rmcp-macros/src/tool.rs index 5e5044e..c640dd5 100644 --- a/crates/rmcp-macros/src/tool.rs +++ b/crates/rmcp-macros/src/tool.rs @@ -7,16 +7,13 @@ use crate::common::extract_doc_line; /// Check if a type is Json and extract the inner type T fn extract_json_inner_type(ty: &syn::Type) -> Option<&syn::Type> { - if let syn::Type::Path(type_path) = ty { - if let Some(last_segment) = type_path.path.segments.last() { - if last_segment.ident == "Json" { - if let syn::PathArguments::AngleBracketed(args) = &last_segment.arguments { - if let Some(syn::GenericArgument::Type(inner_type)) = args.args.first() { - return Some(inner_type); - } - } - } - } + if let syn::Type::Path(type_path) = ty + && let Some(last_segment) = type_path.path.segments.last() + && last_segment.ident == "Json" + && let syn::PathArguments::AngleBracketed(args) = &last_segment.arguments + && let Some(syn::GenericArgument::Type(inner_type)) = args.args.first() + { + return Some(inner_type); } None } @@ -28,13 +25,6 @@ fn extract_schema_from_return_type(ret_type: &syn::Type) -> Option { if let Some(inner_type) = extract_json_inner_type(ret_type) { return syn::parse2::(quote! { rmcp::handler::server::tool::schema_for_output::<#inner_type>() - .unwrap_or_else(|e| { - panic!( - "Invalid output schema for Json<{}>: {}", - std::any::type_name::<#inner_type>(), - e - ) - }) }) .ok(); } @@ -65,13 +55,6 @@ fn extract_schema_from_return_type(ret_type: &syn::Type) -> Option { syn::parse2::(quote! { rmcp::handler::server::tool::schema_for_output::<#inner_type>() - .unwrap_or_else(|e| { - panic!( - "Invalid output schema for Result, E>: {}", - std::any::type_name::<#inner_type>(), - e - ) - }) }) .ok() } @@ -89,8 +72,6 @@ pub struct ToolAttribute { pub output_schema: Option, /// Optional additional tool information. pub annotations: Option, - /// Execution-related configuration including task support. - pub execution: Option, /// Optional icons for the tool pub icons: Option, /// Optional metadata for the tool @@ -100,13 +81,6 @@ pub struct ToolAttribute { pub local: bool, } -#[derive(FromMeta, Debug, Default)] -#[darling(default)] -pub struct ToolExecutionAttribute { - /// Task support mode: "forbidden", "optional", or "required" - pub task_support: Option, -} - pub struct ResolvedToolAttribute { pub name: String, pub title: Option, @@ -114,7 +88,6 @@ pub struct ResolvedToolAttribute { pub input_schema: Expr, pub output_schema: Option, pub annotations: Option, - pub execution: Option, pub icons: Option, pub meta: Option, } @@ -128,7 +101,6 @@ impl ResolvedToolAttribute { input_schema, output_schema, annotations, - execution, icons, meta, } = self; @@ -146,9 +118,6 @@ impl ResolvedToolAttribute { let annotations_call = annotations .map(|a| quote! { .with_annotations(#a) }) .unwrap_or_default(); - let execution_call = execution - .map(|e| quote! { .with_execution(#e) }) - .unwrap_or_default(); let icons_call = icons .map(|i| quote! { .with_icons(#i) }) .unwrap_or_default(); @@ -166,7 +135,6 @@ impl ResolvedToolAttribute { #title_call #output_schema_call #annotations_call - #execution_call #icons_call #meta_call } @@ -278,38 +246,6 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result { } else { None }; - let execution_expr = if let Some(execution) = attribute.execution { - let ToolExecutionAttribute { task_support } = execution; - - let task_support_expr = if let Some(ts) = task_support { - let ts_ident = match ts.as_str() { - "forbidden" => quote! { rmcp::model::TaskSupport::Forbidden }, - "optional" => quote! { rmcp::model::TaskSupport::Optional }, - "required" => quote! { rmcp::model::TaskSupport::Required }, - _ => { - return Err(syn::Error::new( - Span::call_site(), - format!( - "Invalid task_support value '{}'. Expected 'forbidden', 'optional', or 'required'", - ts - ), - )); - } - }; - quote! { Some(#ts_ident) } - } else { - quote! { None } - }; - - let token_stream = quote! { - rmcp::model::ToolExecution::from_raw( - #task_support_expr, - ) - }; - Some(syn::parse2::(token_stream)?) - } else { - None - }; // Handle output_schema - either explicit or generated from return type let output_schema_expr = attribute.output_schema.or_else(|| { // Try to generate schema from return type @@ -333,7 +269,6 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result { input_schema: input_schema_expr, output_schema: output_schema_expr, annotations: annotations_expr, - execution: execution_expr, title: attribute.title, icons: attribute.icons, meta: attribute.meta, @@ -348,13 +283,13 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result { let omit_send = cfg!(feature = "local") || attribute.local; let new_output = syn::parse2::({ let mut lt = quote! { 'static }; - if let Some(receiver) = fn_item.sig.receiver() { - if let Some((_, receiver_lt)) = receiver.reference.as_ref() { - if let Some(receiver_lt) = receiver_lt { - lt = quote! { #receiver_lt }; - } else { - lt = quote! { '_ }; - } + if let Some(receiver) = fn_item.sig.receiver() + && let syn::ReceiverKind::Reference(_, receiver_lt, _) = &receiver.kind + { + if let Some(receiver_lt) = receiver_lt { + lt = quote! { #receiver_lt }; + } else { + lt = quote! { '_ }; } } match &fn_item.sig.output { @@ -407,6 +342,20 @@ mod test { Ok(()) } + #[test] + fn test_async_tool_preserves_receiver_lifetime() -> syn::Result<()> { + let attr = quote! {}; + let input = quote! { + async fn test_tool_with_lifetime<'a>(&'a self) -> String { + "ok".to_string() + } + }; + let result = tool(attr, input)?; + + assert!(result.to_string().contains("+ 'a")); + Ok(()) + } + #[test] fn test_doc_comment_description() -> syn::Result<()> { let attr = quote! {}; // No explicit description diff --git a/crates/rmcp-macros/src/tool_handler.rs b/crates/rmcp-macros/src/tool_handler.rs index 0cb323b..c9f4a25 100644 --- a/crates/rmcp-macros/src/tool_handler.rs +++ b/crates/rmcp-macros/src/tool_handler.rs @@ -47,7 +47,7 @@ pub fn tool_handler(attr: TokenStream, input: TokenStream) -> syn::Result, - ) -> Result { + ) -> Result { let tcc = rmcp::handler::server::tool::ToolCallContext::new(self, request, context); #router.call(tcc).await } @@ -66,12 +66,19 @@ pub fn tool_handler(attr: TokenStream, input: TokenStream) -> syn::Result, - _context: rmcp::service::RequestContext, + context: rmcp::service::RequestContext, ) -> Result { + let supports_cache_hints = context.protocol_version().is_some_and(|version| { + version >= rmcp::model::ProtocolVersion::V_2026_07_28 + }); Ok(rmcp::model::ListToolsResult{ + result_type: Some(rmcp::model::ResultType::COMPLETE), tools: #router.list_all(), meta: #result_meta, next_cursor: None, + ttl_ms: supports_cache_hints.then_some(0), + cache_scope: supports_cache_hints + .then_some(rmcp::model::CacheScope::Public), }) } })?; @@ -107,13 +114,12 @@ pub fn tool_handler(attr: TokenStream, input: TokenStream) -> syn::Result quote! { rmcp::model::Implementation::new(#n, #v) }, (Some(n), None) => { @@ -157,8 +157,8 @@ pub(crate) fn build_get_info( } syn::parse2::(quote! { - fn get_info(&self) -> rmcp::model::ServerInfo { - rmcp::model::ServerInfo::new( + fn get_info(&self) -> rmcp::model::ServerConfig { + rmcp::model::ServerConfig::new( rmcp::model::ServerCapabilities::builder() #(#capability_calls)* .build() diff --git a/crates/rmcp-macros/src/tool_router.rs b/crates/rmcp-macros/src/tool_router.rs index edc8630..6efd692 100644 --- a/crates/rmcp-macros/src/tool_router.rs +++ b/crates/rmcp-macros/src/tool_router.rs @@ -17,6 +17,8 @@ pub struct ToolRouterAttribute { /// When set, also emit `#[::rmcp::tool_handler]` on `impl ServerHandler for Self` so callers /// can skip a separate `#[tool_handler]` block (expanded in a later macro pass). pub server_handler: bool, + /// When set, accept an impl block with no `#[tool]` fn instead of reporting an error. + pub allow_empty: bool, } impl Default for ToolRouterAttribute { @@ -25,6 +27,7 @@ impl Default for ToolRouterAttribute { router: format_ident!("tool_router"), vis: None, server_handler: false, + allow_empty: false, } } } @@ -35,6 +38,7 @@ pub fn tool_router(attr: TokenStream, input: TokenStream) -> syn::Result(input)?; // find all function marked with `#[rmcp::tool]` @@ -58,6 +62,19 @@ pub fn tool_router(attr: TokenStream, input: TokenStream) -> syn::Result syn::Result<()> { + let input = quote! { + impl Probe { + #[tool(description = "probe")] + async fn probe(&self) -> String { "probed".to_owned() } + } + }; + let generated = tool_router(TokenStream::new(), input)?.to_string(); + assert!(generated.contains("with_route"), "{generated}"); + Ok(()) + } + + #[test] + fn tool_router_allow_empty_generates_a_router_without_routes() -> syn::Result<()> { + let generated = tool_router(quote! { allow_empty }, quote! { impl Probe {} })?.to_string(); + assert!(generated.contains("fn tool_router"), "{generated}"); + assert!(!generated.contains("with_route"), "{generated}"); + Ok(()) + } } diff --git a/crates/rmcp/CHANGELOG.md b/crates/rmcp/CHANGELOG.md index e33b77c..a3e4ab7 100644 --- a/crates/rmcp/CHANGELOG.md +++ b/crates/rmcp/CHANGELOG.md @@ -7,6 +7,252 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [3.4.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.3.0...rmcp-v3.4.0) - 2026-09-15 + +### Added + +- *(model)* add ServerConfig and ClientConfig ([#1266](https://github.com/modelcontextprotocol/rust-sdk/pull/1266)) + +### Fixed + +- *(rmcp)* use ServerConfig in cancellation test ([#1273](https://github.com/modelcontextprotocol/rust-sdk/pull/1273)) +- *(rmcp)* route peer cancellation by lifecycle ([#1262](https://github.com/modelcontextprotocol/rust-sdk/pull/1262)) +- run first pre-init request in service loop ([#1263](https://github.com/modelcontextprotocol/rust-sdk/pull/1263)) +- *(auth)* let url origin decide prm discovery ([#1264](https://github.com/modelcontextprotocol/rust-sdk/pull/1264)) +- *(model)* deprecate ServerInfo and ClientInfo aliases ([#1156](https://github.com/modelcontextprotocol/rust-sdk/pull/1156)) +- *(auth)* ignore non-metadata JSON when probing for protected resource metadata ([#1204](https://github.com/modelcontextprotocol/rust-sdk/pull/1204)) +- do not treat malformed JSON 200 as Accepted for requests ([#1208](https://github.com/modelcontextprotocol/rust-sdk/pull/1208)) +- *(streamable-http-server)* map handler-generated HeaderMismatch to HTTP 400 ([#1259](https://github.com/modelcontextprotocol/rust-sdk/pull/1259)) +- *(http)* enforce Origin validation semantics ([#1192](https://github.com/modelcontextprotocol/rust-sdk/pull/1192)) + +## [3.3.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.2.0...rmcp-v3.3.0) - 2026-09-10 + +### Added + +- add ServerHandler::negotiate_initialize ([#1247](https://github.com/modelcontextprotocol/rust-sdk/pull/1247)) +- *(macros)* reject empty tool_router ([#1233](https://github.com/modelcontextprotocol/rust-sdk/pull/1233)) +- *(auth)* add enterprise refresh-token and ID-JAG exchanges ([#1234](https://github.com/modelcontextprotocol/rust-sdk/pull/1234)) + +### Fixed + +- *(sse)* saturate exponential reconnect backoff to avoid overflow panic ([#1231](https://github.com/modelcontextprotocol/rust-sdk/pull/1231)) +- resolve clippy warnings across workspace ([#1195](https://github.com/modelcontextprotocol/rust-sdk/pull/1195)) +- *(auth)* unify refresh checks and error handling ([#1236](https://github.com/modelcontextprotocol/rust-sdk/pull/1236)) + +### Other + +- *(deps)* update process-wrap requirement from 9.0 to 10.0 ([#1229](https://github.com/modelcontextprotocol/rust-sdk/pull/1229)) + +## [3.2.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.1.4...rmcp-v3.2.0) - 2026-08-31 + +### Added + +- *(auth)* coordinate OAuth refreshes through credential stores ([#1232](https://github.com/modelcontextprotocol/rust-sdk/pull/1232)) +- add request-state key rotation ([#1128](https://github.com/modelcontextprotocol/rust-sdk/pull/1128)) + +### Fixed + +- keep initialize on legacy protocol versions ([#1228](https://github.com/modelcontextprotocol/rust-sdk/pull/1228)) +- *(transport)* fall back after sessionless HTTP discover rejections ([#1211](https://github.com/modelcontextprotocol/rust-sdk/pull/1211)) +- allow concurrent streamable http requests ([#1186](https://github.com/modelcontextprotocol/rust-sdk/pull/1186)) + +## [3.1.4](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.1.3...rmcp-v3.1.4) - 2026-08-18 + +### Fixed + +- *(rmcp)* preserve elicitation requestedSchema $schema dialect ([#1176](https://github.com/modelcontextprotocol/rust-sdk/pull/1176)) +- harden signing key handling ([#1166](https://github.com/modelcontextprotocol/rust-sdk/pull/1166)) +- report pre-init metadata errors ([#1160](https://github.com/modelcontextprotocol/rust-sdk/pull/1160)) + +## [3.1.3](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.1.2...rmcp-v3.1.3) - 2026-08-17 + +### Fixed + +- *(auth)* ignore query parameters when matching resources ([#1177](https://github.com/modelcontextprotocol/rust-sdk/pull/1177)) +- *(auth)* retain state until issuer validation ([#1167](https://github.com/modelcontextprotocol/rust-sdk/pull/1167)) +- *(model)* preserve elicitation property order metadata ([#1150](https://github.com/modelcontextprotocol/rust-sdk/pull/1150)) +- time out auto discovery probe ([#1149](https://github.com/modelcontextprotocol/rust-sdk/pull/1149)) +- *(client)* classify discover outcome at source, not at the error type ([#1133](https://github.com/modelcontextprotocol/rust-sdk/pull/1133)) + +### Other + +- Fix typo in to_authorized_http_client doc comment ([#1158](https://github.com/modelcontextprotocol/rust-sdk/pull/1158)) + +## [3.1.2](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.1.1...rmcp-v3.1.2) - 2026-08-07 + +### Fixed + +- *(auth)* map 401/403 challenges on the SSE GET stream ([#1152](https://github.com/modelcontextprotocol/rust-sdk/pull/1152)) +- *(sse)* loop instead of recursing when skipping SSE events ([#1146](https://github.com/modelcontextprotocol/rust-sdk/pull/1146)) +- *(auth)* preserve issuer trailing slash during discovery ([#1145](https://github.com/modelcontextprotocol/rust-sdk/pull/1145)) + +## [3.1.1](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.1.0...rmcp-v3.1.1) - 2026-08-05 + +### Fixed + +- emit cache hints from handler macros ([#1120](https://github.com/modelcontextprotocol/rust-sdk/pull/1120)) +- expose MRTR state to tool handlers ([#1104](https://github.com/modelcontextprotocol/rust-sdk/pull/1104)) +- disambiguate input-required results ([#1103](https://github.com/modelcontextprotocol/rust-sdk/pull/1103)) + +### Other + +- make async-trait optional ([#1119](https://github.com/modelcontextprotocol/rust-sdk/pull/1119)) + +## [3.1.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.0.1...rmcp-v3.1.0) - 2026-07-31 + +### Added + +- classify authorization-required errors ([#1056](https://github.com/modelcontextprotocol/rust-sdk/pull/1056)) +- add strict stateless protocol metadata validation ([#1091](https://github.com/modelcontextprotocol/rust-sdk/pull/1091)) +- SEP-2260 stream-based enforcement of client receive-side request association ([#1055](https://github.com/modelcontextprotocol/rust-sdk/pull/1055)) + +### Fixed + +- *(model)* decode metadata-bearing input-required results affecting mrtr ([#1097](https://github.com/modelcontextprotocol/rust-sdk/pull/1097)) +- require metadata for modern HTTP requests ([#1089](https://github.com/modelcontextprotocol/rust-sdk/pull/1089)) +- honor supported_protocol_versions when negotiating initialize ([#1093](https://github.com/modelcontextprotocol/rust-sdk/pull/1093)) + +### Other + +- document the ping utility with examples ([#1106](https://github.com/modelcontextprotocol/rust-sdk/pull/1106)) +- complete Tier 1 feature docs and finalize roadmap ([#1101](https://github.com/modelcontextprotocol/rust-sdk/pull/1101)) +- *(conformance)* meeting requirements for tier 1 ([#1087](https://github.com/modelcontextprotocol/rust-sdk/pull/1087)) + +## [3.0.1](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.0.0...rmcp-v3.0.1) - 2026-07-29 + +### Fixed + +- *(auth)* use discovered resource for token refresh ([#1084](https://github.com/modelcontextprotocol/rust-sdk/pull/1084)) +- return header mismatch for missing protocol header ([#1083](https://github.com/modelcontextprotocol/rust-sdk/pull/1083)) +- negotiate stateless initialize versions ([#1080](https://github.com/modelcontextprotocol/rust-sdk/pull/1080)) +- stamp server info on graceful subscription results ([#1078](https://github.com/modelcontextprotocol/rust-sdk/pull/1078)) + +## [3.0.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.0.0-beta.5...rmcp-v3.0.0) - 2026-07-28 + +### Fixed + +- recognize 2026 MCP methods ([#1076](https://github.com/modelcontextprotocol/rust-sdk/pull/1076)) + +## [3.0.0-beta.5](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.0.0-beta.4...rmcp-v3.0.0-beta.5) - 2026-07-28 + +### Fixed + +- [**breaking**] remove server_info from DiscoverResult ([#1065](https://github.com/modelcontextprotocol/rust-sdk/pull/1065)) +- preserve transient OAuth discovery HTTP errors ([#1071](https://github.com/modelcontextprotocol/rust-sdk/pull/1071)) +- [**breaking**] preserve OAuth discovery transport errors ([#1069](https://github.com/modelcontextprotocol/rust-sdk/pull/1069)) +- gate client handler bounds for local ([#1068](https://github.com/modelcontextprotocol/rust-sdk/pull/1068)) + +### Other + +- prepare for stable 3.0 release ([#1073](https://github.com/modelcontextprotocol/rust-sdk/pull/1073)) +- RFC 9728 resource is used instead of base url when possible ([#962](https://github.com/modelcontextprotocol/rust-sdk/pull/962)) +- [**breaking**] remove deprecated v3 APIs ([#1066](https://github.com/modelcontextprotocol/rust-sdk/pull/1066)) + +## [3.0.0-beta.4](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.0.0-beta.3...rmcp-v3.0.0-beta.4) - 2026-07-28 + +### Fixed + +- accept namespaced discovery server information ([#1044](https://github.com/modelcontextprotocol/rust-sdk/pull/1044)) + +### Other + +- *(deps)* update base64 requirement from 0.22 to 0.23 ([#1059](https://github.com/modelcontextprotocol/rust-sdk/pull/1059)) +- *(deps)* update jsonwebtoken requirement from 10 to 11 ([#1058](https://github.com/modelcontextprotocol/rust-sdk/pull/1058)) + +## [3.0.0-beta.3](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.0.0-beta.2...rmcp-v3.0.0-beta.3) - 2026-07-27 + +### Fixed + +- [**breaking**] reject missing `issuer` in authorization server metadata by default ([#1054](https://github.com/modelcontextprotocol/rust-sdk/pull/1054)) +- default to allowing missing `issuer` ([#1051](https://github.com/modelcontextprotocol/rust-sdk/pull/1051)) + +### Other + +- make oauth discovery reactive ([#1052](https://github.com/modelcontextprotocol/rust-sdk/pull/1052)) + +## [3.0.0-beta.2](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v3.0.0-beta.1...rmcp-v3.0.0-beta.2) - 2026-07-24 + +### Added + +- expose oauth discovery metadata source ([#1041](https://github.com/modelcontextprotocol/rust-sdk/pull/1041)) + +### Fixed + +- [**breaking**] omit resultType for legacy protocol sessions ([#1038](https://github.com/modelcontextprotocol/rust-sdk/pull/1038)) + +### Other + +- declare and check MSRV ([#1034](https://github.com/modelcontextprotocol/rust-sdk/pull/1034)) + +## [3.0.0-beta.1](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v2.2.0...rmcp-v3.0.0-beta.1) - 2026-07-23 + +### Added + +- route SEP-2260 associated server requests to the originating SSE stream ([#1029](https://github.com/modelcontextprotocol/rust-sdk/pull/1029)) +- [**breaking**] add distributed SSE event store ([#1024](https://github.com/modelcontextprotocol/rust-sdk/pull/1024)) +- add client-side TTL-honoring response cache (SEP-2549) ([#1025](https://github.com/modelcontextprotocol/rust-sdk/pull/1025)) +- [**breaking**] Implement SEP-2663 Tasks Extension ([#1020](https://github.com/modelcontextprotocol/rust-sdk/pull/1020)) +- add subscription listen streams (SEP-2575) ([#1000](https://github.com/modelcontextprotocol/rust-sdk/pull/1000)) +- *(auth)* bind DCR client credentials to issuing authorization server (SEP-2352) ([#998](https://github.com/modelcontextprotocol/rust-sdk/pull/998)) +- add modern client lifecycle modes (SEP-2575) ([#995](https://github.com/modelcontextprotocol/rust-sdk/pull/995)) +- *(auth)* accumulate client-side scopes during step-up authorization ([#888](https://github.com/modelcontextprotocol/rust-sdk/pull/888)) +- [**breaking**] add server discovery and negotiation (SEP-2575) ([#973](https://github.com/modelcontextprotocol/rust-sdk/pull/973)) +- *(conformance)* add SEP-2243 header validation tool ([#997](https://github.com/modelcontextprotocol/rust-sdk/pull/997)) +- [**breaking**] align metadata models with draft schema ([#993](https://github.com/modelcontextprotocol/rust-sdk/pull/993)) +- [**breaking**] implement SEP-2549 cache hints ([#889](https://github.com/modelcontextprotocol/rust-sdk/pull/889)) +- [**breaking**] add MRTR behavior support (SEP-2322) ([#929](https://github.com/modelcontextprotocol/rust-sdk/pull/929)) +- [**breaking**] type Annotations.lastModified as a string ([#956](https://github.com/modelcontextprotocol/rust-sdk/pull/956)) +- [**breaking**] add SEP-2243 HTTP standard headers ([#907](https://github.com/modelcontextprotocol/rust-sdk/pull/907)) +- relax outputSchema to accept non-object JSON Schema types (SEP-2106) ([#895](https://github.com/modelcontextprotocol/rust-sdk/pull/895)) +- [**breaking**] relax tool result structuredContent type (SEP-2106) ([#933](https://github.com/modelcontextprotocol/rust-sdk/pull/933)) +- [**breaking**] add MRTR model types (SEP-2322) ([#915](https://github.com/modelcontextprotocol/rust-sdk/pull/915)) + +### Fixed + +- reap completed response send tasks ([#1026](https://github.com/modelcontextprotocol/rust-sdk/pull/1026)) +- accept stringified numeric response IDs ([#1021](https://github.com/modelcontextprotocol/rust-sdk/pull/1021)) +- re-register after auth server change ([#1011](https://github.com/modelcontextprotocol/rust-sdk/pull/1011)) +- .with_stateful_mode -> .with_legacy_session_mode ([#1015](https://github.com/modelcontextprotocol/rust-sdk/pull/1015)) +- preserve negotiated progress responses ([#1005](https://github.com/modelcontextprotocol/rust-sdk/pull/1005)) +- *(server)* serve draft-version requests statelessly per SEP-2567 ([#999](https://github.com/modelcontextprotocol/rust-sdk/pull/999)) +- pass client header conformance ([#1012](https://github.com/modelcontextprotocol/rust-sdk/pull/1012)) +- *(auth)* add an SDK path for pre-registered OAuth clients ([#994](https://github.com/modelcontextprotocol/rust-sdk/pull/994)) +- *(transport)* cancel in-flight request on stateless streamable-HTTP client disconnect ([#857](https://github.com/modelcontextprotocol/rust-sdk/pull/857)) ([#967](https://github.com/modelcontextprotocol/rust-sdk/pull/967)) +- *(auth)* distinguish rejected refresh tokens from transient failures ([#963](https://github.com/modelcontextprotocol/rust-sdk/pull/963)) +- *(auth)* validate discovered metadata issuer ([#996](https://github.com/modelcontextprotocol/rust-sdk/pull/996)) +- bound streamable HTTP memory usage ([#970](https://github.com/modelcontextprotocol/rust-sdk/pull/970)) +- *(streamable-http)* preserve progress in JSON mode ([#990](https://github.com/modelcontextprotocol/rust-sdk/pull/990)) +- specify compatible sse-stream version ([#968](https://github.com/modelcontextprotocol/rust-sdk/pull/968)) +- flag schema derive on schemars feature ([#966](https://github.com/modelcontextprotocol/rust-sdk/pull/966)) + +### Other + +- refactor OAuth client authorization api ([#1009](https://github.com/modelcontextprotocol/rust-sdk/pull/1009)) +- update for 2026-07-28 version ([#1032](https://github.com/modelcontextprotocol/rust-sdk/pull/1032)) +- *(deps)* update hmac requirement from 0.12 to 0.13 ([#988](https://github.com/modelcontextprotocol/rust-sdk/pull/988)) +- serialize JavaScript dependency install ([#972](https://github.com/modelcontextprotocol/rust-sdk/pull/972)) + +### Changed + +- **BREAKING**: rename `StreamableHttpServerConfig::stateful_mode` to `legacy_session_mode` (and the builder `with_stateful_mode` to `with_legacy_session_mode`) to clarify that the option only affects legacy protocol versions (`< 2026-07-28`); per SEP-2567 the `2026-07-28` draft version is always served statelessly ([#999](https://github.com/modelcontextprotocol/rust-sdk/pull/999)) + +### Fixed + +- use the authorization server issuer as the default `private_key_jwt` audience + +## [2.2.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v2.1.0...rmcp-v2.2.0) - 2026-07-08 + +### Added + +- reject auth servers lacking S256 PKCE support ([#955](https://github.com/modelcontextprotocol/rust-sdk/pull/955)) + +### Fixed + +- pass client conformance suite ([#960](https://github.com/modelcontextprotocol/rust-sdk/pull/960)) +- don't respond to cancelled requests ([#957](https://github.com/modelcontextprotocol/rust-sdk/pull/957)) +- fail orphaned streamable HTTP responses on reinit ([#914](https://github.com/modelcontextprotocol/rust-sdk/pull/914)) +- address 2025-11-25 conformance audit findings ([#951](https://github.com/modelcontextprotocol/rust-sdk/pull/951)) + ## [2.1.0](https://github.com/modelcontextprotocol/rust-sdk/compare/rmcp-v2.0.0...rmcp-v2.1.0) - 2026-07-02 ### Added diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 9704bfc..ac698f9 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -3,6 +3,7 @@ name = "rmcp" license = { workspace = true } version = { workspace = true } edition = { workspace = true } +rust-version = { workspace = true } repository = { workspace = true } homepage = { workspace = true } readme = { workspace = true } @@ -17,11 +18,13 @@ exhaustive_enums = "warn" features = [ "auth", "auth-client-credentials-jwt", + "auth-enterprise-managed", "base64", "client", "client-side-sse", "elicitation", "macros", + "request-state", "reqwest", "reqwest-native-tls", "reqwest-tls-no-provider", @@ -43,12 +46,13 @@ features = [ rustdoc-args = ["--cfg", "docsrs"] [dependencies] -async-trait = "0.1.89" +async-trait = { version = "0.1.89", optional = true } serde = { version = "1.0", features = ["derive", "rc"] } serde_json = "1.0" thiserror = "2" tokio = { version = "1", features = ["sync", "macros", "rt", "time"] } futures = "0.3" +indexmap = { version = "2", features = ["serde"] } tracing = { version = "0.1" } tokio-util = { version = "0.7" } pin-project-lite = "0.2" @@ -56,13 +60,18 @@ pastey = { version = "0.2.0", optional = true } # oauth2 support oauth2 = { version = "5.0", optional = true, default-features = false } # JWT signing for client credentials (private_key_jwt) -jsonwebtoken = { version = "10", optional = true } +jsonwebtoken = { version = "11", optional = true, features = ["aws_lc_rs"] } # for auto generate schema schemars = { version = "1.0", optional = true, features = ["chrono04"] } # for image encoding -base64 = { version = "0.22", optional = true } +base64 = { version = "0.23", optional = true } + +# for SEP-2322 requestState integrity sealing (opt-in via the `request-state` feature) +hmac = { version = "0.13", optional = true } +sha2 = { version = "0.11", optional = true } +zeroize = { version = "1", optional = true } # for HTTP client reqwest = { version = "0.13.2", default-features = false, features = [ @@ -70,7 +79,7 @@ reqwest = { version = "0.13.2", default-features = false, features = [ "stream", ], optional = true } -sse-stream = { version = "0.2", optional = true } +sse-stream = { version = "0.2.4", optional = true } http = { version = "1", optional = true } url = { version = "2.4", optional = true } @@ -79,7 +88,7 @@ url = { version = "2.4", optional = true } tower-service = { version = "0.3", optional = true } # for child process transport -process-wrap = { version = "9.0", features = ["tokio1"], optional = true } +process-wrap = { version = "10.0", features = ["tokio1"], optional = true } # for cross-platform executable path resolution which = { version = "8", optional = true } @@ -116,10 +125,13 @@ chrono = { version = "0.4.38", default-features = false, features = [ default = ["base64", "macros", "server"] local = ["rmcp-macros?/local"] client = ["dep:tokio-stream"] -server = ["transport-async-rw", "dep:schemars", "dep:pastey"] +server = ["transport-async-rw", "schemars", "dep:pastey", "uuid"] macros = ["dep:rmcp-macros", "dep:pastey"] elicitation = ["dep:url"] +# SEP-2322 requestState integrity helper (HMAC-SHA256 seal/open codec) +request-state = ["dep:hmac", "dep:sha2", "dep:zeroize", "base64"] + # reqwest http client __reqwest = ["dep:reqwest"] @@ -139,12 +151,13 @@ server-side-http = [ "dep:bytes", "dep:sse-stream", "tower", + "base64", ] transport-worker = ["dep:tokio-stream"] # SSE stream parsing utilities (used by streamable HTTP client for SSE-formatted responses) -client-side-sse = ["dep:sse-stream", "dep:http"] +client-side-sse = ["dep:sse-stream", "dep:http", "dep:bytes", "base64"] # Streamable HTTP client transport-streamable-http-client = ["client-side-sse", "transport-worker"] @@ -177,16 +190,18 @@ transport-streamable-http-server = [ ] transport-streamable-http-server-session = [ "transport-async-rw", + "dep:async-trait", "dep:tokio-stream", ] # transport-ws = ["transport-io", "dep:tokio-tungstenite"] tower = ["dep:tower-service"] -auth = ["dep:oauth2", "__reqwest", "dep:url"] +auth = ["dep:async-trait", "dep:oauth2", "__reqwest", "dep:url"] auth-client-credentials-jwt = ["auth", "dep:jsonwebtoken", "uuid"] +auth-enterprise-managed = ["auth", "base64"] schemars = ["dep:schemars"] [dev-dependencies] -tokio = { version = "1", features = ["full"] } +tokio = { version = "1", features = ["full", "test-util"] } schemars = { version = "1.1.0", features = ["chrono04"] } axum = { version = "0.8", default-features = false, features = ["http1", "tokio"] } hyper = { version = "1", features = ["server", "http1"] } @@ -200,7 +215,7 @@ tracing-subscriber = { version = "0.3", features = [ "fmt", ] } async-trait = "0.1" -rstest = "0.26.1" +rstest = "0.27.0" [[test]] name = "test_tool_macros" required-features = ["server", "client"] @@ -228,6 +243,11 @@ required-features = [ ] path = "tests/test_with_js.rs" +[[test]] +name = "test_local_client_handler" +required-features = ["client", "local"] +path = "tests/test_local_client_handler.rs" + [[test]] name = "test_notification" required-features = ["server", "client"] @@ -278,6 +298,11 @@ name = "test_streamable_http_json_response" required-features = ["server", "client", "transport-streamable-http-server", "reqwest"] path = "tests/test_streamable_http_json_response.rs" +[[test]] +name = "test_streamable_http_client_concurrency" +required-features = ["client", "transport-streamable-http-client"] +path = "tests/test_streamable_http_client_concurrency.rs" + [[test]] name = "test_streamable_http_protocol_version" required-features = ["server", "client", "transport-streamable-http-server", "reqwest"] @@ -293,11 +318,45 @@ name = "test_protocol_version_negotiation" required-features = ["server", "client"] path = "tests/test_protocol_version_negotiation.rs" +[[test]] +name = "test_client_lifecycle_modes" +required-features = ["client", "server"] +path = "tests/test_client_lifecycle_modes.rs" + +[[test]] +name = "test_stateless_server_requests" +required-features = ["client", "server"] +path = "tests/test_stateless_server_requests.rs" + +[[test]] +name = "test_discover_http_client_startup" +required-features = [ + "client", + "reqwest", + "transport-streamable-http-client-reqwest", + "transport-streamable-http-server", +] +path = "tests/test_discover_http_client_startup.rs" + +[[test]] +name = "test_streamable_http_standard_headers" +required-features = ["server", "client", "transport-streamable-http-server", "reqwest"] +path = "tests/test_streamable_http_standard_headers.rs" + [[test]] name = "test_streamable_http_4xx_error_body" required-features = ["transport-streamable-http-client", "transport-streamable-http-client-reqwest"] path = "tests/test_streamable_http_4xx_error_body.rs" +[[test]] +name = "test_streamable_http_malformed_json_request" +required-features = [ + "client", + "transport-streamable-http-client", + "transport-streamable-http-client-reqwest", +] +path = "tests/test_streamable_http_malformed_json_request.rs" + [[test]] name = "test_custom_request" @@ -309,6 +368,11 @@ name = "test_trace_context" required-features = ["server", "client"] path = "tests/test_trace_context.rs" +[[test]] +name = "test_mrtr_behavior" +required-features = ["server", "client"] +path = "tests/test_mrtr_behavior.rs" + [[test]] name = "test_prompt_macros" required-features = ["server", "client"] @@ -345,6 +409,11 @@ name = "test_client_credentials" required-features = ["auth"] path = "tests/test_client_credentials.rs" +[[test]] +name = "test_auth_error_classification" +required-features = ["auth", "client", "transport-streamable-http-client"] +path = "tests/test_auth_error_classification.rs" + [[test]] name = "test_unix_socket_transport" required-features = [ @@ -357,10 +426,10 @@ path = "tests/test_unix_socket_transport.rs" [[test]] name = "test_streamable_http_stale_session" required-features = [ - "server", - "client", - "transport-streamable-http-server", - "transport-streamable-http-client", + "server", + "client", + "transport-streamable-http-server", + "transport-streamable-http-client", "transport-streamable-http-client-reqwest" ] path = "tests/test_streamable_http_stale_session.rs" @@ -375,6 +444,16 @@ required-features = [ ] path = "tests/test_streamable_http_session_store.rs" +[[test]] +name = "test_streamable_http_event_store" +required-features = [ + "client", + "server", + "transport-streamable-http-client-reqwest", + "transport-streamable-http-server", +] +path = "tests/test_streamable_http_event_store.rs" + [[test]] name = "test_streamable_http_connection_reuse" required-features = [ @@ -387,3 +466,17 @@ required-features = [ "transport-streamable-http-client-reqwest", ] path = "tests/test_streamable_http_connection_reuse.rs" + +[[test]] +name = "test_streamable_http_disconnect_cancel" +required-features = [ + "server", + "transport-streamable-http-server", + "reqwest", +] +path = "tests/test_streamable_http_disconnect_cancel.rs" + +[[test]] +name = "test_sep_2260_stream_routing" +required-features = ["server", "elicitation", "transport-streamable-http-server", "reqwest"] +path = "tests/test_sep_2260_stream_routing.rs" diff --git a/crates/rmcp/README.md b/crates/rmcp/README.md index c133e40..f34469e 100644 --- a/crates/rmcp/README.md +++ b/crates/rmcp/README.md @@ -11,7 +11,7 @@ -The official Rust SDK for the [Model Context Protocol](https://modelcontextprotocol.io/specification/2025-11-25). Build MCP servers that expose tools, resources, and prompts to AI assistants — or build clients that connect to them. +The official Rust SDK for the [Model Context Protocol](https://modelcontextprotocol.io/specification/2026-07-28). Build MCP servers that expose tools, resources, and prompts to AI assistants — or build clients that connect to them. For **getting started**, **usage guides**, and **full MCP feature documentation** (resources, prompts, sampling, roots, logging, completions, subscriptions, etc.), see the [main README](../../README.md). @@ -24,6 +24,7 @@ For **getting started**, **usage guides**, and **full MCP feature documentation* | `macros` | `#[tool]` / `#[prompt]` macros (re-exports [`rmcp-macros`](../rmcp-macros)) | ✅ | | `schemars` | JSON Schema generation for tool definitions | | | `auth` | OAuth 2.0 authentication support | | +| `auth-enterprise-managed` | EMA/XAA refresh-token and ID-JAG exchanges for registered public and confidential clients (includes `auth`) | | | `elicitation` | Elicitation support | | ### Transport features @@ -45,6 +46,10 @@ For **getting started**, **usage guides**, and **full MCP feature documentation* | `reqwest-native-tls` | Uses platform-native TLS (OpenSSL / Secure Transport / SChannel) | | `reqwest-tls-no-provider` | Uses rustls without a default crypto provider (bring your own) | +For enterprise-managed authorization, enable `auth-enterprise-managed` and a TLS +backend such as `reqwest`. See the [EMA/XAA guide](../../docs/OAUTH_SUPPORT.md#enterprise-managed-authorization-emaxaa) +for client authentication and an MCP connection example. + ## Transports The transport layer is pluggable. Two built-in pairs cover the most common cases: diff --git a/crates/rmcp/src/error.rs b/crates/rmcp/src/error.rs index 74f7d43..0940f29 100644 --- a/crates/rmcp/src/error.rs +++ b/crates/rmcp/src/error.rs @@ -1,10 +1,6 @@ use std::{borrow::Cow, fmt::Display}; pub use crate::model::ErrorData; -#[deprecated( - note = "Use `rmcp::ErrorData` instead, `rmcp::ErrorData` could become `RmcpError` in the future." -)] -pub type Error = ErrorData; impl Display for ErrorData { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}: {}", self.code.0, self.message)?; @@ -17,6 +13,20 @@ impl Display for ErrorData { impl std::error::Error for ErrorData {} +#[cfg(all(feature = "auth", any(feature = "client", feature = "server")))] +pub(crate) struct ErrorChain<'a>(pub(crate) &'a (dyn std::error::Error + 'static)); + +#[cfg(all(feature = "auth", any(feature = "client", feature = "server")))] +impl Display for ErrorChain<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0)?; + for source in std::iter::successors(self.0.source(), |source| source.source()) { + write!(f, "\n Caused by: {source}")?; + } + Ok(()) + } +} + /// This is an unified error type for the errors could be returned by the service. #[derive(Debug, thiserror::Error)] #[allow(clippy::large_enum_variant)] diff --git a/crates/rmcp/src/handler/client.rs b/crates/rmcp/src/handler/client.rs index 90f7ed4..486fa3b 100644 --- a/crates/rmcp/src/handler/client.rs +++ b/crates/rmcp/src/handler/client.rs @@ -66,8 +66,8 @@ impl Service for H { ServerNotification::PromptListChangedNotification(_notification_no_param) => { self.on_prompt_list_changed(context).await } - ServerNotification::ElicitationCompleteNotification(notification) => { - self.on_url_elicitation_notification_complete(notification.params, context) + ServerNotification::SubscriptionsAcknowledgedNotification(notification) => { + self.on_subscriptions_acknowledged(notification.params, context) .await } ServerNotification::TaskStatusNotification(notification) => { @@ -85,197 +85,219 @@ impl Service for H { } } -#[allow(unused_variables)] -pub trait ClientHandler: Sized + Send + Sync + 'static { - fn ping( - &self, - context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - std::future::ready(Ok(())) - } +macro_rules! client_handler_methods { + () => { + fn ping( + &self, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + std::future::ready(Ok(())) + } - fn create_message( - &self, - params: CreateMessageRequestParams, - context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - std::future::ready(Err( - McpError::method_not_found::(), - )) - } + fn create_message( + &self, + params: CreateMessageRequestParams, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + std::future::ready(Err( + McpError::method_not_found::(), + )) + } - fn list_roots( - &self, - context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - std::future::ready(Ok(ListRootsResult::default())) - } + fn list_roots( + &self, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + std::future::ready(Ok(ListRootsResult::default())) + } - /// Handle an elicitation request from a server asking for user input. - /// - /// This method is called when a server needs interactive input from the user - /// during tool execution. Implementations should present the message to the user, - /// collect their input according to the requested schema, and return the result. - /// - /// # Arguments - /// * `request` - The elicitation request with message and schema - /// * `context` - The request context - /// - /// # Returns - /// The user's response including action (accept/decline/cancel) and optional data - /// - /// # Default Behavior - /// The default implementation automatically declines all elicitation requests. - /// Real clients should override this to provide user interaction. - /// - /// # Example - /// ```rust,ignore - /// use rmcp::model::ElicitRequestParams; - /// use rmcp::{ - /// model::ErrorData as McpError, - /// model::*, - /// service::{NotificationContext, RequestContext, RoleClient, Service, ServiceRole}, - /// }; - /// use rmcp::ClientHandler; - /// - /// impl ClientHandler for MyClient { - /// async fn create_elicitation( - /// &self, - /// request: ElicitRequestParams, - /// context: RequestContext, - /// ) -> Result { - /// match request { - /// ElicitRequestParams::FormElicitationParam {meta, message, requested_schema,} => { - /// // Display message to user and collect input according to requested_schema - /// let user_input = get_user_input(message, requested_schema).await?; - /// Ok(ElicitResult { - /// action: ElicitationAction::Accept, - /// content: Some(user_input), - /// meta: None, - /// }) - /// } - /// ElicitRequestParams::UrlElicitationParam {meta, message, url, elicitation_id,} => { - /// // Open URL in browser for user to complete elicitation - /// open_url_in_browser(url).await?; - /// Ok(ElicitResult { - /// action: ElicitationAction::Accept, - /// content: None, - /// meta: None, - /// }) - /// } - /// } - /// } - /// } - /// ``` - fn create_elicitation( - &self, - request: ElicitRequestParams, - context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - // Default implementation declines all requests - real clients should override this - let _ = (request, context); - std::future::ready(Ok(ElicitResult { - action: ElicitationAction::Decline, - content: None, - meta: None, - })) - } + /// Handle an elicitation request from a server asking for user input. + /// + /// This method is called when a server needs interactive input from the user + /// during tool execution. Implementations should present the message to the user, + /// collect their input according to the requested schema, and return the result. + /// + /// # Arguments + /// * `request` - The elicitation request with message and schema + /// * `context` - The request context + /// + /// # Returns + /// The user's response including action (accept/decline/cancel) and optional data + /// + /// # Default Behavior + /// The default implementation automatically declines all elicitation requests. + /// Real clients should override this to provide user interaction. + /// + /// # Example + /// ```rust,no_run + /// use rmcp::{ + /// ClientHandler, + /// model::{ + /// ElicitRequestParams, ElicitResult, ElicitationAction, ElicitationSchema, + /// ErrorData as McpError, + /// }, + /// service::{RequestContext, RoleClient}, + /// }; + /// + /// # struct MyClient; + /// # + /// # async fn get_user_input( + /// # _message: String, + /// # _schema: ElicitationSchema, + /// # ) -> Result { + /// # std::future::pending().await + /// # } + /// # + /// # async fn open_url_in_browser(_url: String) -> Result<(), McpError> { + /// # Ok(()) + /// # } + /// # + /// impl ClientHandler for MyClient { + /// async fn create_elicitation( + /// &self, + /// request: ElicitRequestParams, + /// _context: RequestContext, + /// ) -> Result { + /// match request { + /// ElicitRequestParams::FormElicitationParams { + /// message, + /// requested_schema, + /// .. + /// } => { + /// let input = get_user_input(message, requested_schema).await?; + /// Ok(ElicitResult::new(ElicitationAction::Accept).with_content(input)) + /// } + /// ElicitRequestParams::UrlElicitationParams { url, .. } => { + /// open_url_in_browser(url).await?; + /// Ok(ElicitResult::new(ElicitationAction::Accept)) + /// } + /// _ => Ok(ElicitResult::new(ElicitationAction::Decline)), + /// } + /// } + /// } + /// ``` + fn create_elicitation( + &self, + request: ElicitRequestParams, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + // Default implementation declines all requests - real clients should override this + let _ = (request, context); + std::future::ready(Ok(ElicitResult { + action: ElicitationAction::Decline, + content: None, + meta: None, + })) + } - fn on_custom_request( - &self, - request: CustomRequest, - context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - let CustomRequest { method, .. } = request; - let _ = context; - std::future::ready(Err(McpError::new( - ErrorCode::METHOD_NOT_FOUND, - method, - None, - ))) - } + fn on_custom_request( + &self, + request: CustomRequest, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + let CustomRequest { method, .. } = request; + let _ = context; + std::future::ready(Err(McpError::new( + ErrorCode::METHOD_NOT_FOUND, + method, + None, + ))) + } - fn on_cancelled( - &self, - params: CancelledNotificationParam, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - std::future::ready(()) - } - fn on_progress( - &self, - params: ProgressNotificationParam, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - std::future::ready(()) - } - fn on_logging_message( - &self, - params: LoggingMessageNotificationParam, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - std::future::ready(()) - } - fn on_resource_updated( - &self, - params: ResourceUpdatedNotificationParam, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - std::future::ready(()) - } - fn on_resource_list_changed( - &self, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - std::future::ready(()) - } - fn on_tool_list_changed( - &self, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - std::future::ready(()) - } - fn on_prompt_list_changed( - &self, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - std::future::ready(()) - } + fn on_cancelled( + &self, + params: CancelledNotificationParam, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + std::future::ready(()) + } + fn on_progress( + &self, + params: ProgressNotificationParam, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + std::future::ready(()) + } + fn on_logging_message( + &self, + params: LoggingMessageNotificationParam, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + std::future::ready(()) + } + fn on_resource_updated( + &self, + params: ResourceUpdatedNotificationParam, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + std::future::ready(()) + } + fn on_resource_list_changed( + &self, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + std::future::ready(()) + } + fn on_tool_list_changed( + &self, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + std::future::ready(()) + } + fn on_prompt_list_changed( + &self, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + std::future::ready(()) + } + fn on_subscriptions_acknowledged( + &self, + params: SubscriptionsAcknowledgedNotificationParams, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + std::future::ready(()) + } - fn on_url_elicitation_notification_complete( - &self, - params: ElicitationResponseNotificationParam, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - std::future::ready(()) - } - fn on_task_status( - &self, - params: TaskStatusNotificationParam, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - std::future::ready(()) - } - fn on_custom_notification( - &self, - notification: CustomNotification, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - let _ = (notification, context); - std::future::ready(()) - } + fn on_task_status( + &self, + params: TaskStatusNotificationParams, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + std::future::ready(()) + } + fn on_custom_notification( + &self, + notification: CustomNotification, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + let _ = (notification, context); + std::future::ready(()) + } - fn get_info(&self) -> ClientInfo { - ClientInfo::default() - } + fn get_info(&self) -> ClientConfig { + ClientConfig::default() + } + }; } -/// Do nothing, with default client info. +#[allow(unused_variables)] +#[cfg(not(feature = "local"))] +pub trait ClientHandler: Sized + Send + Sync + 'static { + client_handler_methods!(); +} + +#[allow(unused_variables)] +#[cfg(feature = "local")] +pub trait ClientHandler: Sized + 'static { + client_handler_methods!(); +} + +/// Do nothing, with the default client config. impl ClientHandler for () {} -/// Do nothing, with a specific client info. -impl ClientHandler for ClientInfo { - fn get_info(&self) -> ClientInfo { +/// Do nothing, with a specific client config. +impl ClientHandler for ClientConfig { + fn get_info(&self) -> ClientConfig { self.clone() } } @@ -376,9 +398,17 @@ macro_rules! impl_client_handler_for_wrapper { (**self).on_prompt_list_changed(context) } + fn on_subscriptions_acknowledged( + &self, + params: SubscriptionsAcknowledgedNotificationParams, + context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + (**self).on_subscriptions_acknowledged(params, context) + } + fn on_task_status( &self, - params: TaskStatusNotificationParam, + params: TaskStatusNotificationParams, context: NotificationContext, ) -> impl Future + MaybeSendFuture + '_ { (**self).on_task_status(params, context) @@ -392,7 +422,7 @@ macro_rules! impl_client_handler_for_wrapper { (**self).on_custom_notification(notification, context) } - fn get_info(&self) -> ClientInfo { + fn get_info(&self) -> ClientConfig { (**self).get_info() } } diff --git a/crates/rmcp/src/handler/server.rs b/crates/rmcp/src/handler/server.rs index 3cec563..f846724 100644 --- a/crates/rmcp/src/handler/server.rs +++ b/crates/rmcp/src/handler/server.rs @@ -1,13 +1,13 @@ // Sampling/Roots/Logging are SEP-2577-deprecated; internal references are expected. #![expect(deprecated)] -use std::sync::Arc; +use std::{borrow::Cow, sync::Arc}; use crate::{ error::ErrorData as McpError, - model::{TaskSupport, *}, + model::*, service::{ MaybeSendFuture, NotificationContext, RequestContext, RoleServer, Service, ServiceRole, - negotiate_protocol_version, + SubscriptionContext, negotiate_protocol_version, uses_legacy_lifecycle, }, }; @@ -19,6 +19,34 @@ pub mod tool; pub mod tool_name_validation; pub mod wrapper; +/// SEP-2663: gate `tasks/*` methods on the client's declared tasks-extension +/// capability. +/// +/// - If the server does not advertise the tasks extension, the methods are +/// simply unimplemented: `-32601` Method not found. +/// - If the server advertises it but the client did not declare it (either in +/// the request's `_meta` per-request capabilities or, for session-mode +/// peers, at `initialize` time), the spec requires `-32021` Missing +/// Required Client Capability with the required capability in `data`. +fn validate_tasks_capability( + handler: &H, + context: &RequestContext, +) -> Result<(), McpError> { + if !handler.get_info().capabilities.supports_tasks() { + return Err(McpError::method_not_found::()); + } + let client_declared = context + .client_capabilities() + .is_some_and(|caps| caps.supports_tasks()); + if client_declared { + Ok(()) + } else { + Err(McpError::missing_required_client_capability( + ClientCapabilities::builder().enable_tasks().build(), + )) + } +} + impl Service for H { async fn handle_request( &self, @@ -27,13 +55,66 @@ impl Service for H { ) -> Result<::Resp, McpError> { // `context` is moved into the dispatch below, so read the negotiated version first. let protocol_version = context.protocol_version(); + // SEP-2322 (`resultType` discriminator, MRTR) exists from 2026-07-28. + let sep_2322_supported = protocol_version + .as_ref() + .is_some_and(|v| v.as_str() >= ProtocolVersion::V_2026_07_28.as_str()); + let requested_version = context.meta.protocol_version(); + let uses_inline_negotiation = !matches!(&request, ClientRequest::InitializeRequest(_)); + if uses_inline_negotiation && let Some(requested_version) = requested_version.as_ref() { + let supported_versions = self.supported_protocol_versions(); + if !supported_versions.contains(requested_version) { + return Err(McpError::unsupported_protocol_version( + requested_version.clone(), + &supported_versions, + )); + } + } + // Self-contained metadata is required only when the request itself uses + // the inline lifecycle: a discover opener, a session that started without + // `initialize`, or a request that declares 2026-07-28+ in its own _meta. + // Sessions that negotiated via `initialize` (or `serve_directly`) keep the + // session model and may omit per-request metadata. + let requires_request_metadata = uses_inline_negotiation + && (matches!(&request, ClientRequest::DiscoverRequest(_)) + || context.peer.request_metadata_required() + || requested_version.as_ref().is_some_and(|version| { + version.as_str() >= ProtocolVersion::V_2026_07_28.as_str() + })); + if requires_request_metadata { + // Inline lifecycle requests are defined by the 2026-07-28 protocol. + // Validate that lifecycle contract even when a request selects an + // older application protocol version. + let missing = context + .meta + .missing_required_keys(&ProtocolVersion::V_2026_07_28); + if !missing.is_empty() { + return Err(McpError::invalid_params( + format!( + "request _meta is missing or has malformed required fields: {}", + missing.join(", ") + ), + None, + )); + } + } + let legacy_request = + uses_legacy_lifecycle(protocol_version.as_ref(), requires_request_metadata); let result = match request { ClientRequest::InitializeRequest(request) => self .initialize(request.params, context) .await .map(ServerResult::InitializeResult), + ClientRequest::DiscoverRequest(_request) => self + .discover(context) + .await + .map(ServerResult::DiscoverResult), ClientRequest::PingRequest(_request) => { - self.ping(context).await.map(ServerResult::empty) + if !legacy_request { + Err(McpError::method_not_found::()) + } else { + self.ping(context).await.map(ServerResult::empty) + } } ClientRequest::CompleteRequest(request) => self .complete(request.params, context) @@ -46,7 +127,7 @@ impl Service for H { ClientRequest::GetPromptRequest(request) => self .get_prompt(request.params, context) .await - .map(ServerResult::GetPromptResult), + .map(ServerResult::from), ClientRequest::ListPromptsRequest(request) => self .list_prompts(request.params, context) .await @@ -62,52 +143,79 @@ impl Service for H { ClientRequest::ReadResourceRequest(request) => self .read_resource(request.params, context) .await - .map(ServerResult::ReadResourceResult), - ClientRequest::SubscribeRequest(request) => self - .subscribe(request.params, context) - .await - .map(ServerResult::empty), - ClientRequest::UnsubscribeRequest(request) => self - .unsubscribe(request.params, context) - .await - .map(ServerResult::empty), - ClientRequest::CallToolRequest(request) => { - let is_task = request.params.task.is_some(); - - // Validate task support mode per MCP specification - if let Some(tool) = self.get_tool(&request.params.name) { - match (tool.task_support(), is_task) { - // If taskSupport is "required", clients MUST invoke the tool as a task. - // Servers MUST return a -32601 (Method not found) error if they don't. - (TaskSupport::Required, false) => { - return Err(McpError::new( - ErrorCode::METHOD_NOT_FOUND, - "Tool requires task-based invocation", - None, - )); - } - // If taskSupport is "forbidden" (default), clients MUST NOT invoke as a task. - (TaskSupport::Forbidden, true) => { - return Err(McpError::invalid_params( - "Tool does not support task-based invocation", - None, - )); - } - _ => {} - } - } - - if is_task { - tracing::info!("Enqueueing task for tool call: {}", request.params.name); - self.enqueue_task(request.params, context.clone()) - .await - .map(ServerResult::CreateTaskResult) + .map(ServerResult::from), + ClientRequest::SubscriptionsListenRequest(request) => { + if legacy_request { + Err(McpError::method_not_found::()) } else { - self.call_tool(request.params, context) - .await - .map(ServerResult::CallToolResult) + let requested = request.params.notifications; + let Some(candidate) = self.accepted_subscription_filter(&requested) else { + return Err( + McpError::method_not_found::(), + ); + }; + let server_info = self.get_info(); + let advertised = requested.supported_by(&server_info.capabilities); + let handler_accepted = requested.intersection(&candidate); + let accepted = handler_accepted.intersection(&advertised); + if accepted != handler_accepted { + tracing::debug!( + requested_resource_count = requested + .resource_subscriptions + .as_ref() + .map_or(0, Vec::len), + accepted_resource_count = + accepted.resource_subscriptions.as_ref().map_or(0, Vec::len), + "subscription filter reduced to advertised server capabilities" + ); + } + let server_implementation = server_info.server_info; + let subscription_id = context.id.clone(); + let subscription = + SubscriptionContext::establish(context, requested, accepted).await?; + // The 2026-07-28 schema defines a final result for graceful + // server teardown; explicit stdio cancellation remains a notification. + self.listen(subscription).await.map(|()| { + let mut result = SubscriptionsListenResult::complete(subscription_id); + result.meta.set_server_info(server_implementation); + ServerResult::SubscriptionsListenResult(result) + }) } } + ClientRequest::SubscribeRequest(request) => { + if !legacy_request { + Err(McpError::method_not_found::()) + } else { + self.subscribe(request.params, context) + .await + .map(ServerResult::empty) + } + } + ClientRequest::UnsubscribeRequest(request) => { + if !legacy_request { + Err(McpError::method_not_found::()) + } else { + self.unsubscribe(request.params, context) + .await + .map(ServerResult::empty) + } + } + ClientRequest::CallToolRequest(request) => { + let client_declared_tasks = context + .client_capabilities() + .is_some_and(|caps| caps.supports_tasks()); + let response = self.call_tool(request.params, context).await?; + // SEP-2663: the server MUST NOT return CreateTaskResult unless + // the request declared the tasks extension capability. Guard + // against handlers that fail to check before materializing a + // task; such clients cannot parse a task handle. + if matches!(response, CallToolResponse::Task(_)) && !client_declared_tasks { + return Err(McpError::missing_required_client_capability( + ClientCapabilities::builder().enable_tasks().build(), + )); + } + Ok(ServerResult::from(response)) + } ClientRequest::ListToolsRequest(request) => self .list_tools(request.params, context) .await @@ -116,23 +224,41 @@ impl Service for H { .on_custom_request(request, context) .await .map(ServerResult::CustomResult), - ClientRequest::ListTasksRequest(request) => self - .list_tasks(request.params, context) - .await - .map(ServerResult::ListTasksResult), - ClientRequest::GetTaskRequest(request) => self - .get_task_info(request.params, context) - .await - .map(ServerResult::GetTaskResult), - ClientRequest::GetTaskPayloadRequest(request) => self - .get_task_result(request.params, context) - .await - .map(ServerResult::GetTaskPayloadResult), - ClientRequest::CancelTaskRequest(request) => self - .cancel_task(request.params, context) - .await - .map(ServerResult::CancelTaskResult), + ClientRequest::GetTaskRequest(request) => { + validate_tasks_capability::(self, &context)?; + self.get_task(request.params, context) + .await + .map(ServerResult::GetTaskResult) + } + ClientRequest::UpdateTaskRequest(request) => { + validate_tasks_capability::(self, &context)?; + self.update_task(request.params, context) + .await + .map(ServerResult::task_ack) + } + ClientRequest::CancelTaskRequest(request) => { + validate_tasks_capability::(self, &context)?; + self.cancel_task(request.params, context) + .await + .map(ServerResult::task_ack) + } }; + let result = result.and_then(|mut result| { + if matches!(result, ServerResult::InputRequiredResult(_)) && !sep_2322_supported { + Err(McpError::invalid_request( + "InputRequiredResult requires negotiated protocol version 2026-07-28 or newer", + None, + )) + } else { + // Peers on protocol versions older than 2026-07-28 keep the + // legacy wire shape without `resultType: "complete"`. + if !sep_2322_supported { + result.strip_result_type_for_legacy_peer(); + } + Ok(result) + } + }); + // SEP-2164: peers negotiating 2026-07-28+ get the standard INVALID_PARAMS code for // resource-not-found; older peers keep RESOURCE_NOT_FOUND. ISO `YYYY-MM-DD` versions // compare lexically the same as chronologically. @@ -164,9 +290,6 @@ impl Service for H { ClientNotification::RootsListChangedNotification(_notification) => { self.on_roots_list_changed(context).await } - ClientNotification::TaskStatusNotification(notification) => { - self.on_task_status(notification.params, context).await - } ClientNotification::CustomNotification(notification) => { self.on_custom_notification(notification, context).await } @@ -177,20 +300,14 @@ impl Service for H { fn get_info(&self) -> ::Info { self.get_info() } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + ServerHandler::supported_protocol_versions(self) + } } macro_rules! server_handler_methods { () => { - fn enqueue_task( - &self, - _request: CallToolRequestParams, - _context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - std::future::ready(Err(McpError::internal_error( - "Task processing not implemented".to_string(), - None, - ))) - } fn ping( &self, context: RequestContext, @@ -204,12 +321,83 @@ macro_rules! server_handler_methods { context: RequestContext, ) -> impl Future> + MaybeSendFuture + '_ { context.peer.set_peer_info(request.clone()); + std::future::ready(self.negotiate_initialize(&request)) + } + /// Build the `initialize` response for `request`, negotiating the + /// protocol version against [`Self::supported_protocol_versions`]. + /// + /// This is the whole body of the default [`Self::initialize`] minus its + /// `set_peer_info` side effect, so a server that overrides `initialize` + /// to add its own can call this instead of restating the negotiation + /// rule: + /// + /// ``` + /// use rmcp::{ + /// ErrorData as McpError, RoleServer, ServerHandler, + /// model::{InitializeRequestParams, InitializeResult, ServerConfig}, + /// service::RequestContext, + /// }; + /// + /// struct MyServer; + /// + /// impl ServerHandler for MyServer { + /// fn get_info(&self) -> ServerConfig { + /// ServerConfig::default() + /// } + /// + /// async fn initialize( + /// &self, + /// request: InitializeRequestParams, + /// context: RequestContext, + /// ) -> Result { + /// // ... record telemetry, register the peer, etc. + /// context.peer.set_peer_info(request.clone()); + /// self.negotiate_initialize(&request) + /// } + /// } + /// ``` + /// + /// # Errors + /// + /// Returns [`ErrorCode::UNSUPPORTED_PROTOCOL_VERSION`] when this server + /// supports no version that still has an `initialize` handshake. + /// + /// [`ErrorCode::UNSUPPORTED_PROTOCOL_VERSION`]: crate::model::ErrorCode::UNSUPPORTED_PROTOCOL_VERSION + fn negotiate_initialize( + &self, + request: &InitializeRequestParams, + ) -> Result { let mut info = self.get_info(); info.protocol_version = negotiate_protocol_version( &request.protocol_version, - info.protocol_version, - ); - std::future::ready(Ok(info)) + std::mem::take(&mut info.protocol_version), + &self.supported_protocol_versions(), + )?; + Ok(info) + } + /// Return the protocol versions supported by this server. + /// + /// Defaults to every version this SDK knows. Override it to narrow the + /// set to the revisions the server actually implements: the returned + /// list is advertised by [`Self::discover`], bounds what `initialize` + /// negotiation may agree to, and is what per-request versions are + /// validated against. + /// + /// To support everything up to some ceiling, use + /// [`ProtocolVersion::known_up_to`] rather than filtering + /// [`ProtocolVersion::KNOWN_VERSIONS`] by hand. + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(ProtocolVersion::KNOWN_VERSIONS) + } + /// Return this server's discovery information. + fn discover( + &self, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + std::future::ready(Ok(DiscoverResult::from_server_info( + self.supported_protocol_versions().into_owned(), + self.get_info(), + ))) } fn complete( &self, @@ -229,7 +417,7 @@ macro_rules! server_handler_methods { &self, request: GetPromptRequestParams, context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { + ) -> impl Future> + MaybeSendFuture + '_ { std::future::ready(Err(McpError::method_not_found::())) } fn list_prompts( @@ -259,11 +447,41 @@ macro_rules! server_handler_methods { &self, request: ReadResourceRequestParams, context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { + ) -> impl Future> + MaybeSendFuture + '_ { std::future::ready(Err( McpError::method_not_found::(), )) } + /// Return the subset of a requested notification filter this server accepts. + /// + /// Returning `None` leaves `subscriptions/listen` unimplemented. The SDK + /// intersects the returned filter with both `requested` and the notification + /// capabilities advertised by [`Self::get_info`] before acknowledging it. + /// Categories that were not requested or advertised are always removed. + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + None + } + /// Run one established subscription until it is cancelled or closed gracefully. + /// + /// The SDK sends the acknowledgment before invoking this method. Returning + /// `Ok(())` sends the final [`SubscriptionsListenResult`] defined by the + /// 2026-07-28 schema, marking graceful server teardown. Explicit + /// stdio cancellation uses `notifications/cancelled` instead. + fn listen( + &self, + context: SubscriptionContext, + ) -> impl Future> + MaybeSendFuture + '_ { + async move { + context.cancelled().await; + Ok(()) + } + } + #[deprecated( + note = "resources/subscribe is legacy-only; implement accepted_subscription_filter and listen for protocol version 2026-07-28" + )] fn subscribe( &self, request: SubscribeRequestParams, @@ -271,6 +489,9 @@ macro_rules! server_handler_methods { ) -> impl Future> + MaybeSendFuture + '_ { std::future::ready(Err(McpError::method_not_found::())) } + #[deprecated( + note = "resources/unsubscribe is legacy-only; subscriptions/listen is cancelled through its request lifecycle" + )] fn unsubscribe( &self, request: UnsubscribeRequestParams, @@ -312,7 +533,7 @@ macro_rules! server_handler_methods { &self, request: CallToolRequestParams, context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { + ) -> impl Future> + MaybeSendFuture + '_ { std::future::ready(Err(McpError::method_not_found::())) } fn list_tools( @@ -370,13 +591,6 @@ macro_rules! server_handler_methods { ) -> impl Future + MaybeSendFuture + '_ { std::future::ready(()) } - fn on_task_status( - &self, - params: TaskStatusNotificationParam, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - std::future::ready(()) - } fn on_custom_notification( &self, notification: CustomNotification, @@ -386,19 +600,12 @@ macro_rules! server_handler_methods { std::future::ready(()) } - fn get_info(&self) -> ServerInfo { - ServerInfo::default() + fn get_info(&self) -> ServerConfig { + ServerConfig::default() } - fn list_tasks( - &self, - request: Option, - context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - std::future::ready(Err(McpError::method_not_found::())) - } - - fn get_task_info( + /// SEP-2663 `tasks/get`: return the current [`DetailedTask`] state. + fn get_task( &self, request: GetTaskParams, context: RequestContext, @@ -407,20 +614,24 @@ macro_rules! server_handler_methods { std::future::ready(Err(McpError::method_not_found::())) } - fn get_task_result( + /// SEP-2663 `tasks/update`: accept responses to outstanding in-task + /// input requests. Returns an empty acknowledgement on success. + fn update_task( &self, - request: GetTaskPayloadParams, + request: UpdateTaskParams, context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { + ) -> impl Future> + MaybeSendFuture + '_ { let _ = (request, context); - std::future::ready(Err(McpError::method_not_found::())) + std::future::ready(Err(McpError::method_not_found::())) } + /// SEP-2663 `tasks/cancel`: cooperative cancellation. Returns an empty + /// acknowledgement; the task's observable status may lag. fn cancel_task( &self, request: CancelTaskParams, context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { + ) -> impl Future> + MaybeSendFuture + '_ { let _ = (request, context); std::future::ready(Err(McpError::method_not_found::())) } @@ -442,14 +653,6 @@ pub trait ServerHandler: Sized + 'static { macro_rules! impl_server_handler_for_wrapper { ($wrapper:ident) => { impl ServerHandler for $wrapper { - fn enqueue_task( - &self, - request: CallToolRequestParams, - context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - (**self).enqueue_task(request, context) - } - fn ping( &self, context: RequestContext, @@ -465,6 +668,24 @@ macro_rules! impl_server_handler_for_wrapper { (**self).initialize(request, context) } + fn negotiate_initialize( + &self, + request: &InitializeRequestParams, + ) -> Result { + (**self).negotiate_initialize(request) + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + (**self).supported_protocol_versions() + } + + fn discover( + &self, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + (**self).discover(context) + } + fn complete( &self, request: CompleteRequestParams, @@ -485,7 +706,7 @@ macro_rules! impl_server_handler_for_wrapper { &self, request: GetPromptRequestParams, context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { + ) -> impl Future> + MaybeSendFuture + '_ { (**self).get_prompt(request, context) } @@ -518,10 +739,24 @@ macro_rules! impl_server_handler_for_wrapper { &self, request: ReadResourceRequestParams, context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { + ) -> impl Future> + MaybeSendFuture + '_ { (**self).read_resource(request, context) } + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + (**self).accepted_subscription_filter(requested) + } + + fn listen( + &self, + context: SubscriptionContext, + ) -> impl Future> + MaybeSendFuture + '_ { + (**self).listen(context) + } + fn subscribe( &self, request: SubscribeRequestParams, @@ -542,7 +777,7 @@ macro_rules! impl_server_handler_for_wrapper { &self, request: CallToolRequestParams, context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { + ) -> impl Future> + MaybeSendFuture + '_ { (**self).call_tool(request, context) } @@ -596,14 +831,6 @@ macro_rules! impl_server_handler_for_wrapper { (**self).on_roots_list_changed(context) } - fn on_task_status( - &self, - params: TaskStatusNotificationParam, - context: NotificationContext, - ) -> impl Future + MaybeSendFuture + '_ { - (**self).on_task_status(params, context) - } - fn on_custom_notification( &self, notification: CustomNotification, @@ -612,39 +839,31 @@ macro_rules! impl_server_handler_for_wrapper { (**self).on_custom_notification(notification, context) } - fn get_info(&self) -> ServerInfo { + fn get_info(&self) -> ServerConfig { (**self).get_info() } - fn list_tasks( - &self, - request: Option, - context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - (**self).list_tasks(request, context) - } - - fn get_task_info( + fn get_task( &self, request: GetTaskParams, context: RequestContext, ) -> impl Future> + MaybeSendFuture + '_ { - (**self).get_task_info(request, context) + (**self).get_task(request, context) } - fn get_task_result( + fn update_task( &self, - request: GetTaskPayloadParams, + request: UpdateTaskParams, context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - (**self).get_task_result(request, context) + ) -> impl Future> + MaybeSendFuture + '_ { + (**self).update_task(request, context) } fn cancel_task( &self, request: CancelTaskParams, context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { + ) -> impl Future> + MaybeSendFuture + '_ { (**self).cancel_task(request, context) } } diff --git a/crates/rmcp/src/handler/server/common.rs b/crates/rmcp/src/handler/server/common.rs index 0d76547..8d3e4a8 100644 --- a/crates/rmcp/src/handler/server/common.rs +++ b/crates/rmcp/src/handler/server/common.rs @@ -106,32 +106,40 @@ pub fn schema_for_empty_input() -> Arc { EMPTY.clone() } -/// Generate a JSON schema for outputSchema (must have root type "object"; top-level "title" and "description" are removed) -pub fn schema_for_output() -> Result, String> { +/// Strip top-level `title` and `description` from a JSON schema for outputSchema. +/// Unlike `validate_and_strip`, this performs no validation — output schemas are not +/// restricted to `type: "object"` (per SEP-2106). +fn strip_output(raw: &Arc) -> Arc { + let mut object = raw.as_ref().clone(); + object.remove("title"); + object.remove("description"); + Arc::new(object) +} + +/// Generate and strip a JSON schema for outputSchema (top-level "title" and +/// "description" are removed; output schemas are not restricted to root type "object"). +pub fn schema_for_output() -> Arc { thread_local! { - static CACHE_FOR_OUTPUT: std::sync::RwLock, String>>> = Default::default(); + static CACHE_FOR_OUTPUT: std::sync::RwLock>> = Default::default(); }; CACHE_FOR_OUTPUT.with(|cache| { - // Try to get from cache first - if let Some(result) = cache + if let Some(schema) = cache .read() .expect("output schema cache lock poisoned") .get(&TypeId::of::()) { - return result.clone(); + return schema.clone(); } - // Generate, validate, and strip unnecessary top-level fields - let result = validate_and_strip(&schema_for_type::(), "outputSchema"); + let schema = strip_output(&schema_for_type::()); - // Cache the result (both success and error cases) cache .write() .expect("output schema cache lock poisoned") - .insert(TypeId::of::(), result.clone()); + .insert(TypeId::of::(), schema.clone()); - result + schema }) } @@ -201,13 +209,13 @@ where } } -impl FromContextPart for crate::model::Meta +impl FromContextPart for crate::model::RequestMetaObject where C: AsRequestContext, { fn from_context_part(context: &mut C) -> Result { let request_context = context.as_request_context_mut(); - let mut meta = crate::model::Meta::default(); + let mut meta = crate::model::RequestMetaObject::default(); std::mem::swap(&mut meta, &mut request_context.meta); Ok(meta) } @@ -305,10 +313,69 @@ mod tests { assert!(Arc::ptr_eq(&schema, &cloned)); } + #[test] + fn test_schema_for_output_accepts_primitive() { + let schema = schema_for_output::(); + assert_eq!(schema.get("type"), Some(&serde_json::json!("integer"))); + } + + #[test] + fn test_schema_for_output_strips_description_for_primitive() { + let schema = schema_for_output::(); + assert!(!schema.contains_key("description")); + } + + #[test] + fn test_schema_for_output_accepts_composition() { + let schema = schema_for_output::>(); + let schema_str = serde_json::to_string(&schema).unwrap(); + assert!( + schema_str.contains("anyOf") + || schema_str.contains("oneOf") + || schema_str.contains("null"), + "Expected composition schema for Option, got: {schema_str}" + ); + } + + #[test] + fn test_schema_for_output_caches_result() { + let schema1 = schema_for_output::(); + let schema2 = schema_for_output::(); + assert!(Arc::ptr_eq(&schema1, &schema2)); + } + + #[test] + fn test_schema_for_input_rejects_array() { + let result = schema_for_input::>(); + assert!(result.is_err()); + } + + #[test] + fn test_schema_for_output_accepts_unit() { + let _schema = schema_for_output::<()>(); + } + + #[test] + fn test_schema_for_output_accepts_object() { + let schema = schema_for_output::(); + assert_eq!(schema.get("type"), Some(&serde_json::json!("object"))); + } + + #[test] + fn test_schema_for_output_strips_top_level_title() { + let schema = schema_for_output::(); + assert!(!schema.contains_key("title")); + } + + #[test] + fn test_schema_for_output_strips_top_level_description() { + let schema = schema_for_output::(); + assert!(!schema.contains_key("description")); + } + #[rstest] - #[case::output(schema_for_output::)] #[case::input(schema_for_input::)] - fn test_schema_for_object_wrappers_reject_primitives( + fn test_schema_for_input_rejects_primitives( #[case] schema_fn: fn() -> Result, String>, ) { let result = schema_fn(); @@ -316,9 +383,8 @@ mod tests { } #[rstest] - #[case::output(schema_for_output::)] #[case::input(schema_for_input::)] - fn test_schema_for_object_wrappers_accept_objects( + fn test_schema_for_input_accepts_objects( #[case] schema_fn: fn() -> Result, String>, ) { let result = schema_fn(); @@ -326,11 +392,9 @@ mod tests { } #[rstest] - #[case::output_title(schema_for_output::, "title")] - #[case::output_description(schema_for_output::, "description")] #[case::input_title(schema_for_input::, "title")] #[case::input_description(schema_for_input::, "description")] - fn test_schema_for_object_wrappers_strip_top_level_metadata( + fn test_schema_for_input_strips_top_level_metadata( #[case] schema_fn: fn() -> Result, String>, #[case] field: &str, ) { diff --git a/crates/rmcp/src/handler/server/prompt.rs b/crates/rmcp/src/handler/server/prompt.rs index c291ef7..b5b3c46 100644 --- a/crates/rmcp/src/handler/server/prompt.rs +++ b/crates/rmcp/src/handler/server/prompt.rs @@ -15,7 +15,7 @@ pub use super::common::{Extension, RequestId}; use crate::{ RoleServer, handler::server::wrapper::Parameters, - model::{GetPromptResult, PromptMessage}, + model::{GetPromptResponse, GetPromptResult, InputRequiredResult, PromptMessage}, service::{MaybeBoxFuture, MaybeSend, MaybeSendFuture, RequestContext}, }; @@ -59,12 +59,12 @@ pub trait GetPromptHandler { fn handle( self, context: PromptContext<'_, S>, - ) -> MaybeBoxFuture<'_, Result>; + ) -> MaybeBoxFuture<'_, Result>; } /// Type alias for dynamic prompt handlers #[cfg(not(feature = "local"))] -pub type DynGetPromptHandler = dyn for<'a> Fn(PromptContext<'a, S>) -> BoxFuture<'a, Result> +pub type DynGetPromptHandler = dyn for<'a> Fn(PromptContext<'a, S>) -> BoxFuture<'a, Result> + Send + Sync; @@ -73,7 +73,7 @@ pub type DynGetPromptHandler = dyn for<'a> Fn( PromptContext<'a, S>, ) -> futures::future::LocalBoxFuture< 'a, - Result, + Result, >; /// Adapter type for async methods that return `Vec` @@ -91,27 +91,29 @@ pub struct SyncPromptMethodAdapter(PhantomData R>); /// Trait for types that can be converted into GetPromptResult pub trait IntoGetPromptResult { - fn into_get_prompt_result(self) -> Result; + fn into_get_prompt_result(self) -> Result; } impl IntoGetPromptResult for GetPromptResult { - fn into_get_prompt_result(self) -> Result { - Ok(self) + fn into_get_prompt_result(self) -> Result { + Ok(self.into()) + } +} + +impl IntoGetPromptResult for InputRequiredResult { + fn into_get_prompt_result(self) -> Result { + Ok(self.into()) } } impl IntoGetPromptResult for Vec { - fn into_get_prompt_result(self) -> Result { - Ok(GetPromptResult { - description: None, - messages: self, - meta: None, - }) + fn into_get_prompt_result(self) -> Result { + Ok(GetPromptResult::new(self).into()) } } impl IntoGetPromptResult for Result { - fn into_get_prompt_result(self) -> Result { + fn into_get_prompt_result(self) -> Result { self.and_then(|v| v.into_get_prompt_result()) } } @@ -128,7 +130,7 @@ pin_project_lite::pin_project! { }, Ready { #[pin] - result: futures::future::Ready>, + result: futures::future::Ready>, } } } @@ -138,7 +140,7 @@ where F: Future, R: IntoGetPromptResult, { - type Output = Result; + type Output = Result; fn poll( self: std::pin::Pin<&mut Self>, @@ -215,7 +217,7 @@ macro_rules! impl_prompt_handler_for { fn handle( self, mut context: PromptContext<'_, S>, - ) -> MaybeBoxFuture<'_, Result> + ) -> MaybeBoxFuture<'_, Result> { $( let result = $Tn::from_context_part(&mut context); @@ -248,7 +250,7 @@ macro_rules! impl_prompt_handler_for { fn handle( self, mut context: PromptContext<'_, S>, - ) -> MaybeBoxFuture<'_, Result> + ) -> MaybeBoxFuture<'_, Result> { $( let result = $Tn::from_context_part(&mut context); @@ -279,7 +281,7 @@ macro_rules! impl_prompt_handler_for { fn handle( self, mut context: PromptContext<'_, S>, - ) -> MaybeBoxFuture<'_, Result> + ) -> MaybeBoxFuture<'_, Result> { // Extract all parameters before moving into the async block $( @@ -314,7 +316,7 @@ macro_rules! impl_prompt_handler_for { fn handle( self, mut context: PromptContext<'_, S>, - ) -> MaybeBoxFuture<'_, Result> + ) -> MaybeBoxFuture<'_, Result> { $( let result = $Tn::from_context_part(&mut context); diff --git a/crates/rmcp/src/handler/server/router.rs b/crates/rmcp/src/handler/server/router.rs index 45ff9a5..c49c408 100644 --- a/crates/rmcp/src/handler/server/router.rs +++ b/crates/rmcp/src/handler/server/router.rs @@ -1,4 +1,4 @@ -use std::sync::Arc; +use std::{borrow::Cow, sync::Arc}; use prompt::{IntoPromptRoute, PromptRoute}; use tool::{IntoToolRoute, ToolRoute}; @@ -6,7 +6,10 @@ use tool::{IntoToolRoute, ToolRoute}; use super::ServerHandler; use crate::{ RoleServer, Service, - model::{ClientNotification, ClientRequest, ListPromptsResult, ListToolsResult, ServerResult}, + model::{ + ClientNotification, ClientRequest, ListPromptsResult, ListToolsResult, ProtocolVersion, + ServerResult, + }, service::NotificationContext, }; @@ -106,7 +109,7 @@ where context, ); let result = self.tool_router.call(tool_call_context).await?; - Ok(ServerResult::CallToolResult(result)) + Ok(ServerResult::from(result)) } else { self.service .handle_request(ClientRequest::CallToolRequest(request), context) @@ -129,7 +132,7 @@ where context, ); let result = self.prompt_router.get_prompt(prompt_context).await?; - Ok(ServerResult::GetPromptResult(result)) + Ok(ServerResult::from(result)) } else { self.service .handle_request(ClientRequest::GetPromptRequest(request), context) @@ -155,6 +158,10 @@ where .list_changed = Some(true); info } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + ServerHandler::supported_protocol_versions(&self.service) + } } #[cfg(test)] @@ -193,7 +200,7 @@ mod tests { async fn test_router_deferred_notifier_e2e() { let mut router = Router::new(DummyHandler).with_tool(tool::ToolRoute::new_dyn( Tool::new("my_tool", "test", Arc::new(Default::default())), - |_ctx| Box::pin(async { Ok(CallToolResult::default()) }), + |_ctx| Box::pin(async { Ok(CallToolResult::default().into()) }), )); let id_provider: Arc = diff --git a/crates/rmcp/src/handler/server/router/prompt.rs b/crates/rmcp/src/handler/server/router/prompt.rs index e952b2a..509fb42 100644 --- a/crates/rmcp/src/handler/server/router/prompt.rs +++ b/crates/rmcp/src/handler/server/router/prompt.rs @@ -2,7 +2,7 @@ use std::{borrow::Cow, sync::Arc}; use crate::{ handler::server::prompt::{DynGetPromptHandler, GetPromptHandler, PromptContext}, - model::{GetPromptResult, Prompt}, + model::{GetPromptResponse, Prompt}, service::{MaybeBoxFuture, MaybeSend}, }; @@ -50,7 +50,8 @@ impl PromptRoute { where H: for<'a> Fn( PromptContext<'a, S>, - ) -> MaybeBoxFuture<'a, Result> + ) + -> MaybeBoxFuture<'a, Result> + MaybeSend + 'static, { @@ -175,7 +176,7 @@ where pub async fn get_prompt( &self, context: PromptContext<'_, S>, - ) -> Result { + ) -> Result { let item = self.map.get(context.name.as_str()).ok_or_else(|| { crate::ErrorData::invalid_params( format!("prompt '{}' not found", context.name), diff --git a/crates/rmcp/src/handler/server/router/tool.rs b/crates/rmcp/src/handler/server/router/tool.rs index ae096c0..f8fab02 100644 --- a/crates/rmcp/src/handler/server/router/tool.rs +++ b/crates/rmcp/src/handler/server/router/tool.rs @@ -137,21 +137,19 @@ use crate::{ tool::{CallToolHandler, DynCallToolHandler, ToolCallContext}, tool_name_validation::validate_and_warn_tool_name, }, - model::{CallToolResult, ContentBlock, ErrorCode, Tool, ToolAnnotations}, + model::{CallToolResponse, CallToolResult, ContentBlock, ErrorCode, Tool, ToolAnnotations}, service::{MaybeBoxFuture, MaybeSend}, }; const TOOL_ARGUMENT_DESERIALIZATION_ERROR_PREFIX: &str = "failed to deserialize parameters:"; -fn into_tool_argument_error(error: crate::ErrorData) -> Result { +fn into_tool_argument_error(error: crate::ErrorData) -> Result { if error.code == ErrorCode::INVALID_PARAMS && error .message .starts_with(TOOL_ARGUMENT_DESERIALIZATION_ERROR_PREFIX) { - return Ok(CallToolResult::error(vec![ContentBlock::text( - error.message, - )])); + return Ok(CallToolResult::error(vec![ContentBlock::text(error.message)]).into()); } Err(error) @@ -200,7 +198,8 @@ impl ToolRoute { where C: for<'a> Fn( ToolCallContext<'a, S>, - ) -> MaybeBoxFuture<'a, Result> + ) + -> MaybeBoxFuture<'a, Result> + MaybeSend + 'static, { @@ -551,17 +550,17 @@ where } fn notify_if_visible(&self, name: &str) { - if self.map.contains_key(name) { - if let Some(notifier) = &self.notifier { - notifier(); - } + if self.map.contains_key(name) + && let Some(notifier) = &self.notifier + { + notifier(); } } pub async fn call( &self, context: ToolCallContext<'_, S>, - ) -> Result { + ) -> Result { let name = context.name(); if self.disabled.contains(name) { return Err(crate::ErrorData::invalid_params("tool not found", None)); @@ -668,7 +667,8 @@ mod tests { meta: None, name: Cow::Borrowed("requires_params"), arguments: Some(Default::default()), - task: None, + input_responses: None, + request_state: None, }, RequestContext::new(NumberOrString::Number(1), peer), ); @@ -677,6 +677,9 @@ mod tests { .call(ctx) .await .expect("argument validation should be a tool result"); + let CallToolResponse::Complete(result) = result else { + panic!("expected complete CallToolResult"); + }; assert_eq!(result.is_error, Some(true)); let text = result @@ -694,7 +697,7 @@ mod tests { let service = DummyService; let mut router = ToolRouter::new().with_route(ToolRoute::new_dyn( crate::model::Tool::new("test_tool", "a test tool", Arc::new(Default::default())), - |_ctx| Box::pin(async { Ok(CallToolResult::default()) }), + |_ctx| Box::pin(async { Ok(CallToolResult::default().into()) }), )); router.disable_route("test_tool"); @@ -707,7 +710,8 @@ mod tests { meta: None, name: Cow::Borrowed("test_tool"), arguments: None, - task: None, + input_responses: None, + request_state: None, }, RequestContext::new(NumberOrString::Number(1), peer), ); diff --git a/crates/rmcp/src/handler/server/router/tool/tool_traits.rs b/crates/rmcp/src/handler/server/router/tool/tool_traits.rs index b0bf9e2..bde9214 100644 --- a/crates/rmcp/src/handler/server/router/tool/tool_traits.rs +++ b/crates/rmcp/src/handler/server/router/tool/tool_traits.rs @@ -9,7 +9,7 @@ use crate::{ tool::schema_for_output, wrapper::{Json, Parameters}, }, - model::{Icon, JsonObject, Meta, ToolAnnotations, ToolExecution}, + model::{Icon, JsonObject, MetaObject, ToolAnnotations}, schemars::JsonSchema, service::{MaybeSend, MaybeSendFuture}, }; @@ -65,25 +65,16 @@ pub trait ToolBase { /// /// If the tool does not have any output, you should override this methods to return [`None`]. fn output_schema() -> Option> { - Some(schema_for_output::().unwrap_or_else(|e| { - panic!( - "Invalid output schema for ToolBase::Output type `{0}`: {1}", - std::any::type_name::(), - e, - ); - })) + Some(schema_for_output::()) } fn annotations() -> Option { None } - fn execution() -> Option { - None - } fn icons() -> Option> { None } - fn meta() -> Option { + fn meta() -> Option { None } } @@ -117,7 +108,6 @@ pub(crate) fn tool_attribute() -> crate::model::Tool { input_schema: T::input_schema().unwrap_or_else(schema_for_empty_input), output_schema: T::output_schema(), annotations: T::annotations(), - execution: T::execution(), icons: T::icons(), meta: T::meta(), } diff --git a/crates/rmcp/src/handler/server/tool.rs b/crates/rmcp/src/handler/server/tool.rs index 9edf5ff..54de5fa 100644 --- a/crates/rmcp/src/handler/server/tool.rs +++ b/crates/rmcp/src/handler/server/tool.rs @@ -16,7 +16,10 @@ pub use super::{ use crate::{ RoleServer, handler::server::wrapper::Parameters, - model::{CallToolRequestParams, CallToolResult, IntoContents, JsonObject}, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, InputRequiredResult, IntoContents, + JsonObject, + }, service::{MaybeBoxFuture, MaybeSend, MaybeSendFuture, RequestContext}, }; @@ -31,11 +34,18 @@ pub fn parse_json_object(input: JsonObject) -> Result { + /// The request-specific context for this tool call. pub request_context: RequestContext, + /// The server handling this tool call. pub service: &'s S, + /// The name of the tool being called. pub name: Cow<'static, str>, + /// The arguments supplied for the tool call. pub arguments: Option, - pub task: Option, + /// Client responses to input requests from the previous MRTR round. + pub input_responses: Option, + /// Opaque state returned by the server during the previous MRTR round. + pub request_state: Option, } impl<'s, S> ToolCallContext<'s, S> { @@ -45,7 +55,9 @@ impl<'s, S> ToolCallContext<'s, S> { meta: _, name, arguments, - task, + input_responses, + request_state, + .. }: CallToolRequestParams, request_context: RequestContext, ) -> Self { @@ -54,7 +66,8 @@ impl<'s, S> ToolCallContext<'s, S> { service, name, arguments, - task, + input_responses, + request_state, } } pub fn name(&self) -> &str { @@ -76,36 +89,56 @@ impl AsRequestContext for ToolCallContext<'_, S> { } pub trait IntoCallToolResult { - fn into_call_tool_result(self) -> Result; + fn into_call_tool_result(self) -> Result; } impl IntoCallToolResult for T { - fn into_call_tool_result(self) -> Result { - Ok(CallToolResult::success(self.into_contents())) + fn into_call_tool_result(self) -> Result { + Ok(CallToolResult::success(self.into_contents()).into()) } } impl IntoCallToolResult for CallToolResult { - fn into_call_tool_result(self) -> Result { + fn into_call_tool_result(self) -> Result { + Ok(self.into()) + } +} + +impl IntoCallToolResult for InputRequiredResult { + fn into_call_tool_result(self) -> Result { + Ok(self.into()) + } +} + +impl IntoCallToolResult for CallToolResponse { + fn into_call_tool_result(self) -> Result { Ok(self) } } impl IntoCallToolResult for crate::ErrorData { - fn into_call_tool_result(self) -> Result { + fn into_call_tool_result(self) -> Result { Err(self) } } impl IntoCallToolResult for Result { - fn into_call_tool_result(self) -> Result { + fn into_call_tool_result(self) -> Result { match self { Ok(value) => value.into_call_tool_result(), Err(error) => match error.into_call_tool_result() { - Ok(mut result) => { + Ok(CallToolResponse::Complete(mut result)) => { result.is_error = Some(true); - Ok(result) + Ok(result.into()) } + Ok(CallToolResponse::InputRequired(_)) => Err(crate::ErrorData::internal_error( + "InputRequiredResult cannot be returned from a tool error branch", + None, + )), + Ok(CallToolResponse::Task(_)) => Err(crate::ErrorData::internal_error( + "CreateTaskResult cannot be returned from a tool error branch", + None, + )), Err(e) => Err(e), }, } @@ -123,7 +156,7 @@ pin_project_lite::pin_project! { }, Ready { #[pin] - result: Ready>, + result: Ready>, } } } @@ -133,7 +166,7 @@ where F: Future, R: IntoCallToolResult, { - type Output = Result; + type Output = Result; fn poll( self: std::pin::Pin<&mut Self>, @@ -152,20 +185,21 @@ pub trait CallToolHandler { fn call( self, context: ToolCallContext<'_, S>, - ) -> MaybeBoxFuture<'_, Result>; + ) -> MaybeBoxFuture<'_, Result>; } #[cfg(not(feature = "local"))] -pub type DynCallToolHandler = dyn for<'s> Fn(ToolCallContext<'s, S>) -> BoxFuture<'s, Result> +pub type DynCallToolHandler = dyn for<'s> Fn(ToolCallContext<'s, S>) -> BoxFuture<'s, Result> + Send + Sync; #[cfg(feature = "local")] -pub type DynCallToolHandler = - dyn for<'s> Fn( - ToolCallContext<'s, S>, - ) - -> futures::future::LocalBoxFuture<'s, Result>; +pub type DynCallToolHandler = dyn for<'s> Fn( + ToolCallContext<'s, S>, +) -> futures::future::LocalBoxFuture< + 's, + Result, +>; // Tool-specific extractor for tool name #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] @@ -177,6 +211,26 @@ impl FromContextPart> for ToolName { } } +/// Extracts the opaque state returned by the server during the previous MRTR round. +#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] +pub struct RequestState(pub Option); + +impl FromContextPart> for RequestState { + fn from_context_part(context: &mut ToolCallContext) -> Result { + Ok(Self(context.request_state.take())) + } +} + +/// Extracts client responses to input requests from the previous MRTR round. +#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] +pub struct InputResponses(pub Option); + +impl FromContextPart> for InputResponses { + fn from_context_part(context: &mut ToolCallContext) -> Result { + Ok(Self(context.input_responses.take())) + } +} + // Special implementation for Parameters that handles tool arguments impl FromContextPart> for Parameters

where @@ -204,7 +258,10 @@ impl FromContextPart> for JsonObject { } impl<'s, S> ToolCallContext<'s, S> { - pub fn invoke(self, h: H) -> MaybeBoxFuture<'s, Result> + pub fn invoke( + self, + h: H, + ) -> MaybeBoxFuture<'s, Result> where H: CallToolHandler, { @@ -247,7 +304,7 @@ macro_rules! impl_for { fn call( self, mut context: ToolCallContext<'_, S>, - ) -> MaybeBoxFuture<'_, Result>{ + ) -> MaybeBoxFuture<'_, Result>{ $( let result = $Tn::from_context_part(&mut context); let $Tn = match result { @@ -278,7 +335,7 @@ macro_rules! impl_for { fn call( self, mut context: ToolCallContext, - ) -> MaybeBoxFuture<'static, Result>{ + ) -> MaybeBoxFuture<'static, Result>{ $( let result = $Tn::from_context_part(&mut context); let $Tn = match result { @@ -307,7 +364,7 @@ macro_rules! impl_for { fn call( self, mut context: ToolCallContext, - ) -> MaybeBoxFuture<'static, Result> { + ) -> MaybeBoxFuture<'static, Result> { $( let result = $Tn::from_context_part(&mut context); let $Tn = match result { @@ -332,7 +389,7 @@ macro_rules! impl_for { fn call( self, mut context: ToolCallContext, - ) -> MaybeBoxFuture<'static, Result> { + ) -> MaybeBoxFuture<'static, Result> { $( let result = $Tn::from_context_part(&mut context); let $Tn = match result { diff --git a/crates/rmcp/src/handler/server/wrapper/json.rs b/crates/rmcp/src/handler/server/wrapper/json.rs index c03fbd0..7c52979 100644 --- a/crates/rmcp/src/handler/server/wrapper/json.rs +++ b/crates/rmcp/src/handler/server/wrapper/json.rs @@ -3,7 +3,10 @@ use std::borrow::Cow; use schemars::JsonSchema; use serde::Serialize; -use crate::{handler::server::tool::IntoCallToolResult, model::CallToolResult}; +use crate::{ + handler::server::tool::IntoCallToolResult, + model::{CallToolResponse, CallToolResult}, +}; /// Json wrapper for structured output /// @@ -27,7 +30,7 @@ impl JsonSchema for Json { // Implementation for Json to create structured content impl IntoCallToolResult for Json { - fn into_call_tool_result(self) -> Result { + fn into_call_tool_result(self) -> Result { let value = serde_json::to_value(self.0).map_err(|e| { crate::ErrorData::internal_error( format!("Failed to serialize structured content: {}", e), @@ -35,6 +38,6 @@ impl IntoCallToolResult for Json { ) })?; - Ok(CallToolResult::structured(value)) + Ok(CallToolResult::structured(value).into()) } } diff --git a/crates/rmcp/src/lib.rs b/crates/rmcp/src/lib.rs index 9ae3f95..3be6616 100644 --- a/crates/rmcp/src/lib.rs +++ b/crates/rmcp/src/lib.rs @@ -3,8 +3,7 @@ #![doc = include_str!("../README.md")] mod error; -#[allow(deprecated)] -pub use error::{Error, ErrorData, RmcpError}; +pub use error::{ErrorData, RmcpError}; /// Basic data types in MCP specification pub mod model; @@ -16,10 +15,13 @@ pub use handler::client::ClientHandler; pub use handler::server::ServerHandler; #[cfg(feature = "server")] pub use handler::server::wrapper::Json; +#[cfg(feature = "client")] +pub use service::{ + ClientCacheConfig, ClientLifecycleMode, ClientServiceExt, MAX_CLIENT_CACHE_TTL, RoleClient, + select_protocol_version, serve_client, serve_client_with_lifecycle, +}; #[cfg(any(feature = "client", feature = "server"))] pub use service::{Peer, Service, ServiceError, ServiceExt}; -#[cfg(feature = "client")] -pub use service::{RoleClient, serve_client}; #[cfg(feature = "server")] pub use service::{RoleServer, serve_server}; @@ -34,7 +36,7 @@ pub mod transport; pub use pastey::paste; #[cfg(all(feature = "macros", feature = "server"))] pub use rmcp_macros::*; -#[cfg(any(feature = "server", feature = "schemars"))] +#[cfg(feature = "schemars")] pub use schemars; #[cfg(feature = "macros")] pub use serde; diff --git a/crates/rmcp/src/model.rs b/crates/rmcp/src/model.rs index fc45f1e..ec7b93b 100644 --- a/crates/rmcp/src/model.rs +++ b/crates/rmcp/src/model.rs @@ -3,8 +3,10 @@ #![expect(deprecated)] use std::{ borrow::Cow, + collections::hash_map::RandomState, + hash::{BuildHasher, Hasher}, ops::{Deref, DerefMut}, - sync::Arc, + sync::{Arc, OnceLock}, }; mod annotated; mod capabilities; @@ -12,7 +14,10 @@ mod content; mod elicitation_schema; mod extension; mod meta; +mod mrtr; mod prompt; +#[cfg(feature = "request-state")] +mod request_state; mod resource; mod serde_impl; mod task; @@ -23,7 +28,10 @@ pub use content::*; pub use elicitation_schema::*; pub use extension::*; pub use meta::*; +pub use mrtr::*; pub use prompt::*; +#[cfg(feature = "request-state")] +pub use request_state::*; pub use resource::*; use serde::{Deserialize, Serialize, de::DeserializeOwned}; use serde_json::Value; @@ -64,7 +72,7 @@ macro_rules! object { /// without returning any specific data. #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Copy, Eq)] #[serde(deny_unknown_fields)] -#[cfg_attr(feature = "server", derive(schemars::JsonSchema))] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] pub struct EmptyObject {} @@ -166,7 +174,10 @@ impl ProtocolVersion { pub const V_2024_11_05: Self = Self(Cow::Borrowed("2024-11-05")); pub const LATEST: Self = Self::V_2025_11_25; - /// All protocol versions known to this SDK. + /// First protocol version that requires SEP-2243 standard HTTP headers. + pub const STANDARD_HEADERS: Self = Self::V_2026_07_28; + + /// All protocol versions known to this SDK, oldest first. pub const KNOWN_VERSIONS: &[Self] = &[ Self::V_2024_11_05, Self::V_2025_03_26, @@ -179,6 +190,43 @@ impl ProtocolVersion { pub fn as_str(&self) -> &str { &self.0 } + + /// The known versions up to and including `max`, oldest first. + /// + /// Servers that implement every revision up to some ceiling can return + /// this from `supported_protocol_versions` instead of filtering + /// [`Self::KNOWN_VERSIONS`] by hand. `max` itself need not be a known + /// version; the result is empty when it predates all of them. + /// + /// The result borrows from [`Self::KNOWN_VERSIONS`], so call it directly + /// in the method body — it needs no `static` and no `LazyLock`: + /// + /// ```rust,ignore + /// const MAX_SUPPORTED: ProtocolVersion = ProtocolVersion::V_2025_11_25; + /// + /// fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + /// Cow::Borrowed(ProtocolVersion::known_up_to(&MAX_SUPPORTED)) + /// } + /// ``` + /// + /// ``` + /// # use rmcp::model::ProtocolVersion; + /// assert_eq!( + /// ProtocolVersion::known_up_to(&ProtocolVersion::V_2025_06_18), + /// &[ + /// ProtocolVersion::V_2024_11_05, + /// ProtocolVersion::V_2025_03_26, + /// ProtocolVersion::V_2025_06_18, + /// ], + /// ); + /// ``` + pub fn known_up_to(max: &Self) -> &'static [Self] { + let count = Self::KNOWN_VERSIONS + .iter() + .take_while(|version| version.as_str() <= max.as_str()) + .count(); + &Self::KNOWN_VERSIONS[..count] + } } impl Serialize for ProtocolVersion { @@ -229,6 +277,22 @@ impl NumberOrString { NumberOrString::String(s) => Value::String(s.to_string()), } } + + pub(crate) fn numeric_string_value(&self) -> Option { + match self { + Self::String(id) => id.parse().ok(), + Self::Number(_) => None, + } + } + + pub(crate) fn matches_response_id(&self, response_id: &Self) -> bool { + self == response_id + || matches!( + self, + Self::Number(request_id) + if response_id.numeric_string_value() == Some(*request_id) + ) + } } impl std::fmt::Display for NumberOrString { @@ -330,7 +394,8 @@ pub struct ProgressToken(pub NumberOrString); pub struct Request { pub method: M, pub params: P, - /// extensions will carry anything possible in the context, including [`Meta`] + /// extensions will carry anything possible in the context, including the metadata + /// ([`RequestMetaObject`] for requests, [`NotificationMetaObject`] for notifications) /// /// this is similar with the Extensions in `http` crate #[cfg_attr(feature = "schemars", schemars(skip))] @@ -363,7 +428,8 @@ pub struct RequestOptionalParam { pub method: M, // #[serde(skip_serializing_if = "Option::is_none")] pub params: Option

, - /// extensions will carry anything possible in the context, including [`Meta`] + /// extensions will carry anything possible in the context, including the metadata + /// ([`RequestMetaObject`] for requests, [`NotificationMetaObject`] for notifications) /// /// this is similar with the Extensions in `http` crate #[cfg_attr(feature = "schemars", schemars(skip))] @@ -385,7 +451,8 @@ impl RequestOptionalParam { #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] pub struct RequestNoParam { pub method: M, - /// extensions will carry anything possible in the context, including [`Meta`] + /// extensions will carry anything possible in the context, including the metadata + /// ([`RequestMetaObject`] for requests, [`NotificationMetaObject`] for notifications) /// /// this is similar with the Extensions in `http` crate #[cfg_attr(feature = "schemars", schemars(skip))] @@ -406,7 +473,8 @@ impl GetExtensions for RequestNoParam { pub struct Notification { pub method: M, pub params: P, - /// extensions will carry anything possible in the context, including [`Meta`] + /// extensions will carry anything possible in the context, including the metadata + /// ([`RequestMetaObject`] for requests, [`NotificationMetaObject`] for notifications) /// /// this is similar with the Extensions in `http` crate #[cfg_attr(feature = "schemars", schemars(skip))] @@ -428,7 +496,8 @@ impl Notification { #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] pub struct NotificationNoParam { pub method: M, - /// extensions will carry anything possible in the context, including [`Meta`] + /// extensions will carry anything possible in the context, including the metadata + /// ([`RequestMetaObject`] for requests, [`NotificationMetaObject`] for notifications) /// /// this is similar with the Extensions in `http` crate #[cfg_attr(feature = "schemars", schemars(skip))] @@ -471,9 +540,9 @@ pub struct JsonRpcResponse { #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] pub struct JsonRpcError { pub jsonrpc: JsonRpcVersion2_0, - // MCP 2025-11-25 §Error Responses: `id` is optional and omitted when the + // MCP 2026-07-28 §Error Responses: `id` is optional and omitted when the // server cannot read the request id (e.g. parse error / invalid request). - // https://modelcontextprotocol.io/specification/2025-11-25/basic#error-responses + // https://modelcontextprotocol.io/specification/2026-07-28/basic#error-responses #[serde(default, skip_serializing_if = "Option::is_none")] pub id: Option, pub error: ErrorData, @@ -510,13 +579,17 @@ pub struct JsonRpcNotification { pub struct ErrorCode(pub i32); impl ErrorCode { + /// The request used a protocol version the server does not support. + pub const UNSUPPORTED_PROTOCOL_VERSION: Self = Self(-32022); + /// Processing the request requires a client capability that was not declared. + pub const MISSING_REQUIRED_CLIENT_CAPABILITY: Self = Self(-32021); + pub const HEADER_MISMATCH: Self = Self(-32020); pub const RESOURCE_NOT_FOUND: Self = Self(-32002); pub const INVALID_REQUEST: Self = Self(-32600); pub const METHOD_NOT_FOUND: Self = Self(-32601); pub const INVALID_PARAMS: Self = Self(-32602); pub const INTERNAL_ERROR: Self = Self(-32603); pub const PARSE_ERROR: Self = Self(-32700); - pub const URL_ELICITATION_REQUIRED: Self = Self(-32042); } /// Error information for JSON-RPC error responses. @@ -540,6 +613,8 @@ pub struct ErrorData { } impl ErrorData { + const TRANSPORT_CLOSED_MARKER: &str = "io.modelcontextprotocol/transportClosed"; + pub fn new( code: ErrorCode, message: impl Into>, @@ -556,7 +631,33 @@ impl ErrorData { pub fn resource_not_found(message: impl Into>, data: Option) -> Self { Self::new(ErrorCode::RESOURCE_NOT_FOUND, message, data) } - + pub fn header_mismatch(message: impl Into>, data: Option) -> Self { + Self::new(ErrorCode::HEADER_MISMATCH, message, data) + } + /// Create an unsupported-protocol-version error. + pub fn unsupported_protocol_version( + requested: ProtocolVersion, + supported: &[ProtocolVersion], + ) -> Self { + Self::new( + ErrorCode::UNSUPPORTED_PROTOCOL_VERSION, + "Unsupported protocol version", + Some(serde_json::json!({ + "requested": requested, + "supported": supported, + })), + ) + } + /// Create a missing-required-capability error. + pub fn missing_required_client_capability(required: ClientCapabilities) -> Self { + Self::new( + ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY, + "Missing required client capability", + Some(serde_json::json!({ + "requiredCapabilities": required, + })), + ) + } pub fn parse_error(message: impl Into>, data: Option) -> Self { Self::new(ErrorCode::PARSE_ERROR, message, data) } @@ -572,11 +673,32 @@ impl ErrorData { pub fn internal_error(message: impl Into>, data: Option) -> Self { Self::new(ErrorCode::INTERNAL_ERROR, message, data) } - pub fn url_elicitation_required( - message: impl Into>, - data: Option, - ) -> Self { - Self::new(ErrorCode::URL_ELICITATION_REQUIRED, message, data) + + #[cfg(feature = "transport-streamable-http-client")] + pub(crate) fn transport_closed(message: impl Into>) -> Self { + let mut data = JsonObject::new(); + data.insert( + Self::TRANSPORT_CLOSED_MARKER.to_owned(), + Value::from(Self::transport_closed_token()), + ); + Self::internal_error(message, Some(Value::Object(data))) + } + + pub(crate) fn is_transport_closed(&self) -> bool { + self.data + .as_ref() + .and_then(|data| data.get(Self::TRANSPORT_CLOSED_MARKER)) + .and_then(Value::as_u64) + == Some(Self::transport_closed_token()) + } + + fn transport_closed_token() -> u64 { + static TOKEN: OnceLock = OnceLock::new(); + *TOKEN.get_or_init(|| { + let mut hasher = RandomState::new().build_hasher(); + hasher.write(b"rmcp transport-closed marker"); + hasher.finish() + }) } } @@ -684,6 +806,85 @@ impl From for () { fn from(_value: EmptyResult) {} } +/// Indicates the type of a result object, allowing the client to +/// determine how to parse the response. +/// +/// The spec defines this as an open string (`"complete" | "input_required" | string`), +/// so unknown values are preserved rather than rejected. Servers implementing this +/// protocol version MUST include `resultType` in every result. For backward +/// compatibility, clients MUST treat an absent field as `"complete"`. +/// +/// Ordinary results model the field as `Option`: `None` means the +/// field is absent on the wire. Constructors default to `Some(COMPLETE)`, and +/// the server handler strips the `"complete"` discriminator before responding +/// to peers that negotiated a protocol version older than `2026-07-28`, so +/// legacy sessions keep their historical wire shape (see +/// [`ServerResult::strip_result_type_for_legacy_peer`]). +#[derive(Debug, Clone, PartialEq, Eq)] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +pub struct ResultType(Cow<'static, str>); + +impl ResultType { + pub const COMPLETE: Self = Self(Cow::Borrowed("complete")); + pub const INPUT_REQUIRED: Self = Self(Cow::Borrowed("input_required")); + /// SEP-2663 Tasks extension: the result is a task handle ([`CreateTaskResult`]). + pub const TASK: Self = Self(Cow::Borrowed("task")); + + pub fn as_str(&self) -> &str { + &self.0 + } + + /// Returns `true` if this is `"input_required"`. + pub fn is_input_required(&self) -> bool { + self.0 == "input_required" + } + + /// Returns `true` if this is `"complete"`. + pub fn is_complete(&self) -> bool { + self.0 == "complete" + } + + /// Returns `true` if this is `"task"` (SEP-2663 Tasks extension). + pub fn is_task(&self) -> bool { + self.0 == "task" + } +} + +impl Default for ResultType { + fn default() -> Self { + Self::COMPLETE + } +} + +impl Serialize for ResultType { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + self.0.serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for ResultType { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let s: String = Deserialize::deserialize(deserializer)?; + match s.as_str() { + "complete" => Ok(Self::COMPLETE), + "input_required" => Ok(Self::INPUT_REQUIRED), + _ => Ok(Self(Cow::Owned(s))), + } + } +} + +impl std::fmt::Display for ResultType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + /// A catch-all response either side can use for custom requests. #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] #[serde(transparent)] @@ -712,7 +913,7 @@ pub struct CancelledNotificationParam { #[serde(skip_serializing_if = "Option::is_none")] pub reason: Option, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } impl CancelledNotificationParam { @@ -748,7 +949,8 @@ pub type CancelledNotification = pub struct CustomNotification { pub method: String, pub params: Option, - /// extensions will carry anything possible in the context, including [`Meta`] + /// extensions will carry anything possible in the context, including the metadata + /// ([`RequestMetaObject`] for requests, [`NotificationMetaObject`] for notifications) /// /// this is similar with the Extensions in `http` crate #[cfg_attr(feature = "schemars", schemars(skip))] @@ -783,7 +985,8 @@ impl CustomNotification { pub struct CustomRequest { pub method: String, pub params: Option, - /// extensions will carry anything possible in the context, including [`Meta`] + /// extensions will carry anything possible in the context, including the metadata + /// ([`RequestMetaObject`] for requests, [`NotificationMetaObject`] for notifications) /// /// this is similar with the Extensions in `http` crate #[cfg_attr(feature = "schemars", schemars(skip))] @@ -828,7 +1031,7 @@ pub type InitializedNotification = NotificationNoParam, + pub meta: Option, /// The MCP protocol version this client supports pub protocol_version: ProtocolVersion, /// The capabilities this client supports (sampling, roots, etc.) @@ -855,18 +1058,14 @@ impl InitializeRequestParams { } impl RequestParamsMeta for InitializeRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -/// Deprecated: Use [`InitializeRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use InitializeRequestParams instead")] -pub type InitializeRequestParam = InitializeRequestParams; - /// The server's response to an initialization request. /// /// Contains the server's protocol version, capabilities, and implementation @@ -886,7 +1085,7 @@ pub struct InitializeResult { #[serde(skip_serializing_if = "Option::is_none")] pub instructions: Option, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } impl InitializeResult { @@ -920,13 +1119,273 @@ impl InitializeResult { } } +/// A server's own configuration: the protocol version, capabilities, +/// implementation identity, and instructions it advertises to clients. +/// +/// This is what `ServerHandler::get_info` returns. It is an alias for +/// [`InitializeResult`] because the same value is sent as the `initialize` +/// response on the wire. +/// +/// # Examples +/// +/// ``` +/// use rmcp::model::{ServerCapabilities, ServerConfig}; +/// +/// let config = ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) +/// .with_instructions("Call `add` to sum two numbers."); +/// assert!(config.capabilities.tools.is_some()); +/// ``` +pub type ServerConfig = InitializeResult; + +/// A client's own configuration: the protocol version, capabilities, and +/// implementation identity it advertises to servers. +/// +/// This is what `ClientHandler::get_info` returns. It is an alias for +/// [`InitializeRequestParams`] because the same value is sent as the +/// `initialize` request on the wire. +/// +/// # Examples +/// +/// ``` +/// use rmcp::model::{ClientCapabilities, ClientConfig, Implementation}; +/// +/// let config = ClientConfig::new( +/// ClientCapabilities::builder().enable_elicitation().build(), +/// Implementation::new("my-client", "1.0.0"), +/// ); +/// assert!(config.capabilities.elicitation.is_some()); +/// ``` +pub type ClientConfig = InitializeRequestParams; + +/// Deprecated alias for [`ServerConfig`]. +/// +/// The name collides with the protocol's `serverInfo` field, which is only the +/// [`Implementation`] identity, so `server_info.server_info` was easy to +/// misread (#1082). +#[deprecated(note = "use `ServerConfig` instead")] pub type ServerInfo = InitializeResult; + +/// Deprecated alias for [`ClientConfig`]. +/// +/// The name collides with the protocol's `clientInfo` field, which is only the +/// [`Implementation`] identity, so `client_info.client_info` was easy to +/// misread (#1082). +#[deprecated(note = "use `ClientConfig` instead")] pub type ClientInfo = InitializeRequestParams; +/// Information negotiated about a server peer. +/// +/// Unlike [`InitializeResult`], the server implementation identity is optional +/// because discovery responses are not required to provide it. +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] +#[serde(rename_all = "camelCase")] +#[non_exhaustive] +pub struct ServerPeerInfo { + /// The negotiated MCP protocol version. + pub protocol_version: ProtocolVersion, + /// The capabilities this server provides. + pub capabilities: ServerCapabilities, + /// Information about the server implementation, when provided. + #[serde(skip_serializing_if = "Option::is_none")] + pub server_info: Option, + /// Optional human-readable instructions about using this server. + #[serde(skip_serializing_if = "Option::is_none")] + pub instructions: Option, + /// Protocol-level response metadata. + #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] + pub meta: Option, +} + +impl ServerPeerInfo { + /// Create peer information without a server implementation identity. + pub fn new(protocol_version: ProtocolVersion, capabilities: ServerCapabilities) -> Self { + Self { + protocol_version, + capabilities, + server_info: None, + instructions: None, + meta: None, + } + } + + /// Set the server implementation identity. + pub fn with_server_info(mut self, server_info: Implementation) -> Self { + self.server_info = Some(server_info); + self + } + + /// Set instructions supplied by the server. + pub fn with_instructions(mut self, instructions: impl Into) -> Self { + self.instructions = Some(instructions.into()); + self + } +} + +impl From for ServerPeerInfo { + fn from(result: InitializeResult) -> Self { + Self { + protocol_version: result.protocol_version, + capabilities: result.capabilities, + server_info: Some(result.server_info), + instructions: result.instructions, + meta: result.meta, + } + } +} + +const_string!(DiscoverRequestMethod = "server/discover"); + +/// Parameters for [`DiscoverRequest`]. +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, Default)] +#[serde(deny_unknown_fields)] +#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] +pub struct DiscoverRequestParams {} + +#[cfg(feature = "schemars")] +#[derive(schemars::JsonSchema)] +#[expect(dead_code, reason = "schema-only representation of request parameters")] +struct DiscoverRequestParamsSchema { + #[schemars(rename = "_meta")] + meta: RequestMetaObject, +} + +#[cfg(feature = "schemars")] +impl schemars::JsonSchema for DiscoverRequestParams { + fn schema_name() -> Cow<'static, str> { + Cow::Borrowed("DiscoverRequestParams") + } + + fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema { + DiscoverRequestParamsSchema::json_schema(generator) + } +} + +/// A request for the server's supported protocol versions and capabilities. +pub type DiscoverRequest = Request; + +/// The server's response to a [`DiscoverRequest`]. +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[serde(rename_all = "camelCase")] +#[non_exhaustive] +pub struct DiscoverResult { + /// Identifies how the result should be parsed. + pub result_type: ResultType, + /// Protocol versions implemented by this server. + pub supported_versions: Vec, + /// Capabilities provided by this server. + pub capabilities: ServerCapabilities, + /// Optional guidance for using the server. + #[serde(skip_serializing_if = "Option::is_none")] + pub instructions: Option, + /// How long clients may consider this response fresh, in milliseconds. + pub ttl_ms: u64, + /// Whether the cached result may be shared across authorization contexts. + pub cache_scope: CacheScope, + /// Protocol-level response metadata. + #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] + pub meta: Option, +} + +const SERVER_INFO_META_KEY: &str = "io.modelcontextprotocol/serverInfo"; + +fn server_info_from_meta(meta: &MetaObject) -> Option { + meta.get(SERVER_INFO_META_KEY) + .and_then(|value| serde_json::from_value(value.clone()).ok()) +} + +fn set_server_info_on_meta(meta: &mut MetaObject, server_info: Implementation) { + let server_info = + serde_json::to_value(server_info).expect("Implementation serialization cannot fail"); + meta.insert(SERVER_INFO_META_KEY.to_owned(), server_info); +} + +impl DiscoverResult { + /// Create a non-cacheable private discovery result. + pub fn new(supported_versions: Vec, capabilities: ServerCapabilities) -> Self { + Self { + result_type: ResultType::COMPLETE, + supported_versions, + capabilities, + instructions: None, + ttl_ms: 0, + cache_scope: CacheScope::Private, + meta: None, + } + } + + /// Return the server implementation information stored in result metadata. + pub fn server_info(&self) -> Option { + server_info_from_meta(self.meta.as_ref()?) + } + + /// Store server implementation information in result metadata. + pub fn set_server_info(&mut self, server_info: Implementation) { + set_server_info_on_meta(self.meta.get_or_insert_default(), server_info); + } + + /// Store server implementation information in result metadata. + pub fn with_server_info(mut self, server_info: Implementation) -> Self { + self.set_server_info(server_info); + self + } + + /// Create a discovery result from the server's configuration. + pub fn from_server_info( + supported_versions: Vec, + server_config: ServerConfig, + ) -> Self { + let ServerConfig { + capabilities, + server_info, + instructions, + meta, + .. + } = server_config; + let mut result = Self { + result_type: ResultType::COMPLETE, + supported_versions, + capabilities, + instructions, + ttl_ms: 0, + cache_scope: CacheScope::Private, + meta, + }; + result.set_server_info(server_info); + result + } + + /// Set the cache lifetime hint in milliseconds. + pub fn with_ttl_ms(mut self, ttl_ms: u64) -> Self { + self.ttl_ms = ttl_ms; + self + } + + /// Set the cache scope. + pub fn with_cache_scope(mut self, cache_scope: CacheScope) -> Self { + self.cache_scope = cache_scope; + self + } +} + +impl ServerPeerInfo { + /// Create peer information from a discovery result and the selected version. + pub fn from_discover_result(protocol_version: ProtocolVersion, result: DiscoverResult) -> Self { + let server_info = result.server_info(); + Self { + protocol_version, + capabilities: result.capabilities, + server_info, + instructions: result.instructions, + meta: result.meta, + } + } +} + #[allow(clippy::derivable_impls)] -impl Default for ServerInfo { +impl Default for InitializeResult { fn default() -> Self { - ServerInfo { + InitializeResult { protocol_version: ProtocolVersion::default(), capabilities: ServerCapabilities::default(), server_info: Implementation::from_build_env(), @@ -937,9 +1396,9 @@ impl Default for ServerInfo { } #[allow(clippy::derivable_impls)] -impl Default for ClientInfo { +impl Default for InitializeRequestParams { fn default() -> Self { - ClientInfo { + InitializeRequestParams { meta: None, protocol_version: ProtocolVersion::default(), capabilities: ClientCapabilities::default(), @@ -1097,7 +1556,7 @@ impl Implementation { pub struct PaginatedRequestParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, #[serde(skip_serializing_if = "Option::is_none")] pub cursor: Option, } @@ -1110,17 +1569,14 @@ impl PaginatedRequestParams { } impl RequestParamsMeta for PaginatedRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -/// Deprecated: Use [`PaginatedRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use PaginatedRequestParams instead")] -pub type PaginatedRequestParam = PaginatedRequestParams; // ============================================================================= // PROGRESS AND PAGINATION // ============================================================================= @@ -1144,7 +1600,7 @@ pub struct ProgressNotificationParam { #[serde(skip_serializing_if = "Option::is_none")] pub message: Option, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } impl ProgressNotificationParam { @@ -1176,32 +1632,105 @@ pub type ProgressNotification = Notification= 0`; if a server returns a negative value, +/// clients SHOULD treat it as `0` (immediately stale). This tolerates that case +/// rather than erroring, while still accepting an absent field as `None`. +fn deserialize_ttl_ms<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let value = Option::::deserialize(deserializer)?; + Ok(value.map(|ttl_ms| ttl_ms.max(0) as u64)) +} + macro_rules! paginated_result { ($t:ident { $i_item: ident: $t_item: ty }) => { - #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] + #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] pub struct $t { - #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, - #[serde(skip_serializing_if = "Option::is_none")] + /// Result type discriminator (SEP-2322). Required by the [spec schema] + /// for servers implementing protocol version `2026-07-28`, but optional + /// here because this type also models results from older protocol + /// versions, which do not carry the field: `None` means absent on the + /// wire, and per the spec "the client MUST treat the absent field as + /// `"complete"`". Constructors default to `Some(ResultType::COMPLETE)`; + /// the server handler clears the field when responding to peers that + /// negotiated an older version. + /// + /// [spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub result_type: Option, + #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] + pub meta: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub next_cursor: Option, + /// Time, in milliseconds, that this result may be treated as fresh (SEP-2549). + /// Required by spec version 2026-07-28, but optional here to maintain compatibility + /// with older spec versions. + #[serde( + default, + deserialize_with = "deserialize_ttl_ms", + skip_serializing_if = "Option::is_none" + )] + pub ttl_ms: Option, + /// Scope describing who may cache this result (SEP-2549). + /// Required by spec version 2026-07-28, but optional here to maintain compatibility + /// with older spec versions. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cache_scope: Option, pub $i_item: $t_item, } + impl Default for $t { + fn default() -> Self { + Self::with_all_items(Default::default()) + } + } + impl $t { - pub fn with_all_items( - items: $t_item, - ) -> Self { + pub fn with_all_items(items: $t_item) -> Self { Self { + result_type: Some(ResultType::COMPLETE), meta: None, next_cursor: None, + ttl_ms: None, + cache_scope: None, $i_item: items, } } + + /// Set the time, in milliseconds, that this result may be treated as fresh. + pub fn with_ttl_ms(mut self, ttl_ms: u64) -> Self { + self.ttl_ms = Some(ttl_ms); + self + } + + /// Set the cache scope for this result. + pub fn with_cache_scope(mut self, cache_scope: CacheScope) -> Self { + self.cache_scope = Some(cache_scope); + self + } } }; } @@ -1237,9 +1766,16 @@ const_string!(ReadResourceRequestMethod = "resources/read"); pub struct ReadResourceRequestParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// The URI of the resource to read pub uri: String, + /// Client responses to server-initiated input requests from a previous + /// [`InputRequiredResult`]. + #[serde(skip_serializing_if = "Option::is_none")] + pub input_responses: Option, + /// Opaque request state echoed back from a previous [`InputRequiredResult`]. + #[serde(skip_serializing_if = "Option::is_none")] + pub request_state: Option, } impl ReadResourceRequestParams { @@ -1248,48 +1784,100 @@ impl ReadResourceRequestParams { Self { meta: None, uri: uri.into(), + input_responses: None, + request_state: None, } } /// Set the metadata for this request. - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: RequestMetaObject) -> Self { self.meta = Some(meta); self } + + /// Sets the input responses for an MRTR retry. + pub fn with_input_responses(mut self, input_responses: InputResponses) -> Self { + self.input_responses = Some(input_responses); + self + } + + /// Sets the request state for an MRTR retry. + pub fn with_request_state(mut self, request_state: impl Into) -> Self { + self.request_state = Some(request_state.into()); + self + } } impl RequestParamsMeta for ReadResourceRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -/// Deprecated: Use [`ReadResourceRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use ReadResourceRequestParams instead")] -pub type ReadResourceRequestParam = ReadResourceRequestParams; - /// Result containing the contents of a read resource #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] +#[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub struct ReadResourceResult { + /// Result type discriminator (SEP-2322). Required by the [spec schema] + /// for servers implementing protocol version `2026-07-28`, but optional + /// here because this type also models results from older protocol + /// versions, which do not carry the field: `None` means absent on the + /// wire, and per the spec "the client MUST treat the absent field as + /// `"complete"`". Constructors default to `Some(ResultType::COMPLETE)`; + /// the server handler clears the field when responding to peers that + /// negotiated an older version. + /// + /// [spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub result_type: Option, + /// Time, in milliseconds, that this result may be treated as fresh (SEP-2549). + /// Required by spec version 2026-07-28, but optional here to maintain compatibility + /// with older spec versions. + #[serde( + default, + deserialize_with = "deserialize_ttl_ms", + skip_serializing_if = "Option::is_none" + )] + pub ttl_ms: Option, + /// Scope describing who may cache this result (SEP-2549). + /// Required by spec version 2026-07-28, but optional here to maintain compatibility + /// with older spec versions. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cache_scope: Option, /// The actual content of the resource pub contents: Vec, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } impl ReadResourceResult { /// Create a new ReadResourceResult with the given contents. pub fn new(contents: Vec) -> Self { Self { + result_type: Some(ResultType::COMPLETE), + ttl_ms: None, + cache_scope: None, contents, meta: None, } } + + /// Set the time, in milliseconds, that this result may be treated as fresh. + pub fn with_ttl_ms(mut self, ttl_ms: u64) -> Self { + self.ttl_ms = Some(ttl_ms); + self + } + + /// Set the cache scope for this result. + pub fn with_cache_scope(mut self, cache_scope: CacheScope) -> Self { + self.cache_scope = Some(cache_scope); + self + } } /// Request to read a specific resource @@ -1309,7 +1897,7 @@ const_string!(SubscribeRequestMethod = "resources/subscribe"); pub struct SubscribeRequestParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// The URI of the resource to subscribe to pub uri: String, } @@ -1325,19 +1913,18 @@ impl SubscribeRequestParams { } impl RequestParamsMeta for SubscribeRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -/// Deprecated: Use [`SubscribeRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use SubscribeRequestParams instead")] -pub type SubscribeRequestParam = SubscribeRequestParams; - /// Request to subscribe to resource updates +#[deprecated( + note = "resources/subscribe is legacy-only; use subscriptions/listen for protocol version 2026-07-28" +)] pub type SubscribeRequest = Request; const_string!(UnsubscribeRequestMethod = "resources/unsubscribe"); @@ -1349,7 +1936,7 @@ const_string!(UnsubscribeRequestMethod = "resources/unsubscribe"); pub struct UnsubscribeRequestParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// The URI of the resource to unsubscribe from pub uri: String, } @@ -1365,19 +1952,18 @@ impl UnsubscribeRequestParams { } impl RequestParamsMeta for UnsubscribeRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -/// Deprecated: Use [`UnsubscribeRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use UnsubscribeRequestParams instead")] -pub type UnsubscribeRequestParam = UnsubscribeRequestParams; - /// Request to unsubscribe from resource updates +#[deprecated( + note = "resources/unsubscribe is legacy-only; cancel the subscriptions/listen request for protocol version 2026-07-28" +)] pub type UnsubscribeRequest = Request; const_string!(ResourceUpdatedNotificationMethod = "notifications/resources/updated"); @@ -1390,7 +1976,7 @@ pub struct ResourceUpdatedNotificationParam { /// The URI of the resource that was updated pub uri: String, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } impl ResourceUpdatedNotificationParam { @@ -1407,6 +1993,405 @@ impl ResourceUpdatedNotificationParam { pub type ResourceUpdatedNotification = Notification; +// ============================================================================= +// SUBSCRIPTIONS +// ============================================================================= + +/// Notification categories a client opts in to on a `subscriptions/listen` stream. +#[derive(Debug, Default, Serialize, Deserialize, Clone, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[non_exhaustive] +pub struct SubscriptionFilter { + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "schemars", schemars(with = "bool"))] + pub tools_list_changed: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "schemars", schemars(with = "bool"))] + pub prompts_list_changed: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "schemars", schemars(with = "bool"))] + pub resources_list_changed: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "schemars", schemars(with = "Vec"))] + pub resource_subscriptions: Option>, +} + +impl SubscriptionFilter { + /// Create an empty filter that opts in to no notifications. + pub fn new() -> Self { + Self::default() + } + + /// Create a builder for a subscription filter. + pub fn builder() -> SubscriptionFilterBuilder { + SubscriptionFilterBuilder::default() + } + + /// Return the subset present in both filters. + pub fn intersection(&self, other: &Self) -> Self { + let resource_subscriptions = self + .resource_subscriptions + .as_ref() + .and_then(|requested| { + other.resource_subscriptions.as_ref().map(|accepted| { + requested + .iter() + .filter(|uri| accepted.contains(uri)) + .cloned() + .collect() + }) + }) + .filter(|uris: &Vec| !uris.is_empty()); + Self { + tools_list_changed: (self.tools_list_changed == Some(true) + && other.tools_list_changed == Some(true)) + .then_some(true), + prompts_list_changed: (self.prompts_list_changed == Some(true) + && other.prompts_list_changed == Some(true)) + .then_some(true), + resources_list_changed: (self.resources_list_changed == Some(true) + && other.resources_list_changed == Some(true)) + .then_some(true), + resource_subscriptions, + } + } + + /// Return whether this filter accepts only notifications requested by `other`. + pub fn is_subset_of(&self, other: &Self) -> bool { + let booleans_are_subset = [ + (self.tools_list_changed, other.tools_list_changed), + (self.prompts_list_changed, other.prompts_list_changed), + (self.resources_list_changed, other.resources_list_changed), + ] + .into_iter() + .all(|(accepted, requested)| accepted != Some(true) || requested == Some(true)); + let resources_are_subset = self.resource_subscriptions.as_ref().is_none_or(|accepted| { + accepted.iter().all(|uri| { + other + .resource_subscriptions + .as_ref() + .is_some_and(|requested| requested.contains(uri)) + }) + }); + booleans_are_subset && resources_are_subset + } + + /// Return the requested notification types advertised by server capabilities. + pub fn supported_by(&self, capabilities: &ServerCapabilities) -> Self { + Self { + tools_list_changed: (self.tools_list_changed == Some(true) + && capabilities + .tools + .as_ref() + .is_some_and(|tools| tools.list_changed == Some(true))) + .then_some(true), + prompts_list_changed: (self.prompts_list_changed == Some(true) + && capabilities + .prompts + .as_ref() + .is_some_and(|prompts| prompts.list_changed == Some(true))) + .then_some(true), + resources_list_changed: (self.resources_list_changed == Some(true) + && capabilities + .resources + .as_ref() + .is_some_and(|resources| resources.list_changed == Some(true))) + .then_some(true), + resource_subscriptions: capabilities + .resources + .as_ref() + .is_some_and(|resources| resources.subscribe == Some(true)) + .then(|| self.resource_subscriptions.clone()) + .flatten(), + } + } +} + +/// Builder for [`SubscriptionFilter`]. +#[derive(Debug, Default)] +#[non_exhaustive] +pub struct SubscriptionFilterBuilder { + filter: SubscriptionFilter, +} + +impl SubscriptionFilterBuilder { + /// Opt in to `notifications/tools/list_changed`. + pub fn tools_list_changed(mut self) -> Self { + self.filter.tools_list_changed = Some(true); + self + } + + /// Opt in to `notifications/prompts/list_changed`. + pub fn prompts_list_changed(mut self) -> Self { + self.filter.prompts_list_changed = Some(true); + self + } + + /// Opt in to `notifications/resources/list_changed`. + pub fn resources_list_changed(mut self) -> Self { + self.filter.resources_list_changed = Some(true); + self + } + + /// Opt in to updates for all supplied resource URIs. + pub fn resource_subscriptions( + mut self, + uris: impl IntoIterator>, + ) -> Self { + self.filter.resource_subscriptions = Some(uris.into_iter().map(Into::into).collect()); + self + } + + /// Add one resource URI to the update subscription set. + pub fn resource_subscription(mut self, uri: impl Into) -> Self { + self.filter + .resource_subscriptions + .get_or_insert_default() + .push(uri.into()); + self + } + + /// Build the filter. + pub fn build(self) -> SubscriptionFilter { + self.filter + } +} + +const_string!(SubscriptionsListenRequestMethod = "subscriptions/listen"); + +#[cfg(feature = "schemars")] +fn subscriptions_listen_request_meta_schema( + generator: &mut schemars::SchemaGenerator, +) -> schemars::Schema { + let progress_token = generator.subschema_for::(); + let client_info = generator.subschema_for::(); + let client_capabilities = generator.subschema_for::(); + let log_level = generator.subschema_for::(); + schemars::json_schema!({ + "type": "object", + "properties": { + "progressToken": progress_token, + "io.modelcontextprotocol/protocolVersion": { + "type": "string", + }, + "io.modelcontextprotocol/clientInfo": client_info, + "io.modelcontextprotocol/clientCapabilities": client_capabilities, + "io.modelcontextprotocol/logLevel": log_level, + }, + "required": RequestMetaObject::DRAFT_REQUIRED_KEYS, + "additionalProperties": true, + }) +} + +/// Parameters for opening a long-lived notification subscription. +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] +#[serde(rename_all = "camelCase")] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[non_exhaustive] +pub struct SubscriptionsListenRequestParams { + /// Protocol-level metadata. Required by the 2026-07-28 wire schema. + #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] + #[cfg_attr( + feature = "schemars", + schemars(required, schema_with = "subscriptions_listen_request_meta_schema") + )] + pub meta: Option, + /// Notification categories requested for this stream. + pub notifications: SubscriptionFilter, +} + +impl SubscriptionsListenRequestParams { + /// Create listen parameters for a notification filter. + pub fn new(notifications: SubscriptionFilter) -> Self { + Self { + meta: None, + notifications, + } + } + + /// Set protocol-level request metadata. + pub fn with_meta(mut self, meta: RequestMetaObject) -> Self { + self.meta = Some(meta); + self + } +} + +impl RequestParamsMeta for SubscriptionsListenRequestParams { + fn meta(&self) -> Option<&RequestMetaObject> { + self.meta.as_ref() + } + + fn meta_mut(&mut self) -> &mut Option { + &mut self.meta + } +} + +/// Request that opens a long-lived notification subscription. +pub type SubscriptionsListenRequest = + Request; + +const SUBSCRIPTION_ID_META_KEY: &str = "io.modelcontextprotocol/subscriptionId"; + +/// Metadata on the final result of a `subscriptions/listen` request. +#[derive(Debug, Serialize, Clone, PartialEq)] +#[serde(transparent)] +#[non_exhaustive] +pub struct SubscriptionsListenResultMeta(MetaObject); + +impl SubscriptionsListenResultMeta { + /// Create result metadata for the originating listen request. + pub fn new(subscription_id: RequestId) -> Self { + let mut meta = MetaObject::new(); + meta.insert( + SUBSCRIPTION_ID_META_KEY.to_owned(), + subscription_id.into_json_value(), + ); + Self(meta) + } + + /// Return the originating listen request ID, if the metadata remains valid. + pub fn subscription_id(&self) -> Option { + self.0 + .get(SUBSCRIPTION_ID_META_KEY) + .and_then(|value| RequestId::deserialize(value).ok()) + } + + /// Replace the originating listen request ID. + pub fn set_subscription_id(&mut self, subscription_id: RequestId) { + self.0.insert( + SUBSCRIPTION_ID_META_KEY.to_owned(), + subscription_id.into_json_value(), + ); + } + + /// Return the server implementation information stored in result metadata. + pub fn server_info(&self) -> Option { + server_info_from_meta(&self.0) + } + + /// Store server implementation information in result metadata. + pub fn set_server_info(&mut self, server_info: Implementation) { + set_server_info_on_meta(&mut self.0, server_info); + } +} + +impl<'de> Deserialize<'de> for SubscriptionsListenResultMeta { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let meta = MetaObject::deserialize(deserializer)?; + let Some(value) = meta.get(SUBSCRIPTION_ID_META_KEY) else { + return Err(serde::de::Error::missing_field(SUBSCRIPTION_ID_META_KEY)); + }; + RequestId::deserialize(value).map_err(serde::de::Error::custom)?; + Ok(Self(meta)) + } +} + +impl std::ops::Deref for SubscriptionsListenResultMeta { + type Target = MetaObject; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl std::ops::DerefMut for SubscriptionsListenResultMeta { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +#[cfg(feature = "schemars")] +impl schemars::JsonSchema for SubscriptionsListenResultMeta { + fn schema_name() -> Cow<'static, str> { + Cow::Borrowed("SubscriptionsListenResultMeta") + } + + fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema { + let subscription_id = generator.subschema_for::(); + let server_info = generator.subschema_for::(); + schemars::json_schema!({ + "type": "object", + "properties": { + "io.modelcontextprotocol/serverInfo": { + "description": "Identifies the server software producing the response. Servers SHOULD include this field on every response unless specifically configured not to do so.", + "allOf": [server_info], + }, + "io.modelcontextprotocol/subscriptionId": subscription_id, + }, + "required": ["io.modelcontextprotocol/subscriptionId"], + "additionalProperties": true, + }) + } +} + +/// Final response indicating that a subscription ended gracefully. +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] +#[serde(rename_all = "camelCase")] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[non_exhaustive] +pub struct SubscriptionsListenResult { + pub result_type: ResultType, + #[serde(rename = "_meta")] + pub meta: SubscriptionsListenResultMeta, +} + +impl SubscriptionsListenResult { + /// Create a completed subscription result. + pub fn new(meta: SubscriptionsListenResultMeta) -> Self { + Self { + result_type: ResultType::COMPLETE, + meta, + } + } + + /// Create a completed result for the originating listen request. + pub fn complete(subscription_id: RequestId) -> Self { + Self::new(SubscriptionsListenResultMeta::new(subscription_id)) + } +} + +const_string!( + SubscriptionsAcknowledgedNotificationMethod = "notifications/subscriptions/acknowledged" +); + +/// Parameters reporting the accepted subset of a subscription filter. +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] +#[serde(rename_all = "camelCase")] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[non_exhaustive] +pub struct SubscriptionsAcknowledgedNotificationParams { + #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "schemars", schemars(with = "NotificationMetaObject"))] + pub meta: Option, + pub notifications: SubscriptionFilter, +} + +impl SubscriptionsAcknowledgedNotificationParams { + /// Create acknowledgment parameters for the accepted filter. + pub fn new(notifications: SubscriptionFilter) -> Self { + Self { + meta: None, + notifications, + } + } + + /// Set notification metadata. + pub fn with_meta(mut self, meta: NotificationMetaObject) -> Self { + self.meta = Some(meta); + self + } +} + +/// First notification sent on an established subscription stream. +pub type SubscriptionsAcknowledgedNotification = Notification< + SubscriptionsAcknowledgedNotificationMethod, + SubscriptionsAcknowledgedNotificationParams, +>; + // ============================================================================= // PROMPT MANAGEMENT // ============================================================================= @@ -1429,10 +2414,17 @@ const_string!(GetPromptRequestMethod = "prompts/get"); pub struct GetPromptRequestParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, pub name: String, #[serde(skip_serializing_if = "Option::is_none")] pub arguments: Option, + /// Client responses to server-initiated input requests from a previous + /// [`InputRequiredResult`]. + #[serde(skip_serializing_if = "Option::is_none")] + pub input_responses: Option, + /// Opaque request state echoed back from a previous [`InputRequiredResult`]. + #[serde(skip_serializing_if = "Option::is_none")] + pub request_state: Option, } impl GetPromptRequestParams { @@ -1442,6 +2434,8 @@ impl GetPromptRequestParams { meta: None, name: name.into(), arguments: None, + input_responses: None, + request_state: None, } } @@ -1452,25 +2446,33 @@ impl GetPromptRequestParams { } /// Set the metadata for this request. - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: RequestMetaObject) -> Self { self.meta = Some(meta); self } + + /// Sets the input responses for an MRTR retry. + pub fn with_input_responses(mut self, input_responses: InputResponses) -> Self { + self.input_responses = Some(input_responses); + self + } + + /// Sets the request state for an MRTR retry. + pub fn with_request_state(mut self, request_state: impl Into) -> Self { + self.request_state = Some(request_state.into()); + self + } } impl RequestParamsMeta for GetPromptRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -/// Deprecated: Use [`GetPromptRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use GetPromptRequestParams instead")] -pub type GetPromptRequestParam = GetPromptRequestParams; - /// Request to get a specific prompt pub type GetPromptRequest = Request; @@ -1519,7 +2521,7 @@ const_string!(SetLevelRequestMethod = "logging/setLevel"); pub struct SetLevelRequestParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// The desired logging level pub level: LoggingLevel, } @@ -1532,18 +2534,14 @@ impl SetLevelRequestParams { } impl RequestParamsMeta for SetLevelRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -/// Deprecated: Use [`SetLevelRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use SetLevelRequestParams instead")] -pub type SetLevelRequestParam = SetLevelRequestParams; - /// Request to set the logging level #[deprecated( since = "2.0.0", @@ -1570,7 +2568,7 @@ pub struct LoggingMessageNotificationParam { /// The actual log data pub data: Value, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } impl LoggingMessageNotificationParam { @@ -1791,7 +2789,7 @@ pub struct SamplingMessage { /// The actual content of the message (text, image, audio, tool use, or tool result) pub content: SamplingContent, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } /// Content types for sampling messages (SEP-1577). @@ -1813,9 +2811,6 @@ pub enum SamplingMessageContentBlock { ToolResult(ToolResultContent), } -#[deprecated(since = "2.0.0", note = "Renamed to SamplingMessageContentBlock")] -pub type SamplingMessageContent = SamplingMessageContentBlock; - impl SamplingMessageContentBlock { /// Create a text content pub fn text(text: impl Into) -> Self { @@ -1958,10 +2953,7 @@ pub enum ContextInclusion { pub struct CreateMessageRequestParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, - /// Task metadata for async task management (SEP-1319) - #[serde(skip_serializing_if = "Option::is_none")] - pub task: Option, + pub meta: Option, /// The conversation history and current messages pub messages: Vec, /// Preferences for model selection and behavior @@ -1993,29 +2985,19 @@ pub struct CreateMessageRequestParams { } impl RequestParamsMeta for CreateMessageRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -impl TaskAugmentedRequestParamsMeta for CreateMessageRequestParams { - fn task(&self) -> Option<&TaskMetadata> { - self.task.as_ref() - } - fn task_mut(&mut self) -> &mut Option { - &mut self.task - } -} - impl CreateMessageRequestParams { /// Create a new CreateMessageRequestParams with required fields. pub fn new(messages: Vec, max_tokens: u32) -> Self { Self { meta: None, - task: None, messages, model_preferences: None, system_prompt: None, @@ -2155,10 +3137,6 @@ impl CreateMessageRequestParams { } } -/// Deprecated: Use [`CreateMessageRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use CreateMessageRequestParams instead")] -pub type CreateMessageRequestParam = CreateMessageRequestParams; - /// Preferences for model selection and behavior in sampling requests. /// /// This allows servers to express their preferences for which model to use @@ -2313,7 +3291,7 @@ impl CompletionContext { pub struct CompleteRequestParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, pub r#ref: Reference, pub argument: ArgumentInfo, /// Optional context containing previously resolved argument values @@ -2340,18 +3318,14 @@ impl CompleteRequestParams { } impl RequestParamsMeta for CompleteRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -/// Deprecated: Use [`CompleteRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use CompleteRequestParams instead")] -pub type CompleteRequestParam = CompleteRequestParams; - pub type CompleteRequest = Request; #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] @@ -2433,20 +3407,39 @@ impl CompletionInfo { } } -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub struct CompleteResult { + /// Result type discriminator (SEP-2322). Required by the [spec schema] + /// for servers implementing protocol version `2026-07-28`, but optional + /// here because this type also models results from older protocol + /// versions, which do not carry the field: `None` means absent on the + /// wire, and per the spec "the client MUST treat the absent field as + /// `"complete"`". Constructors default to `Some(ResultType::COMPLETE)`; + /// the server handler clears the field when responding to peers that + /// negotiated an older version. + /// + /// [spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub result_type: Option, pub completion: CompletionInfo, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, +} + +impl Default for CompleteResult { + fn default() -> Self { + Self::new(CompletionInfo::default()) + } } impl CompleteResult { /// Create a new CompleteResult with the given completion info. pub fn new(completion: CompletionInfo) -> Self { Self { + result_type: Some(ResultType::COMPLETE), completion, meta: None, } @@ -2519,9 +3512,6 @@ impl ResourceTemplateReference { } } -#[deprecated(since = "2.0.0", note = "Renamed to ResourceTemplateReference")] -pub type ResourceReference = ResourceTemplateReference; - #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] @@ -2582,7 +3572,7 @@ pub struct Root { #[serde(skip_serializing_if = "Option::is_none")] pub name: Option, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } impl Root { @@ -2602,7 +3592,7 @@ impl Root { } /// Sets the protocol-level metadata for this root. - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } @@ -2626,7 +3616,7 @@ pub type ListRootsRequest = RequestNoParam; pub struct ListRootsResult { pub roots: Vec, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } impl ListRootsResult { @@ -2636,7 +3626,7 @@ impl ListRootsResult { } /// Sets the protocol-level metadata for this result. - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } @@ -2653,7 +3643,6 @@ pub type RootsListChangedNotification = NotificationNoParam, + meta: Option, message: String, requested_schema: ElicitationSchema, }, #[serde(rename = "url", rename_all = "camelCase")] - UrlElicitationParam { + Url { #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - meta: Option, + meta: Option, message: String, url: String, elicitation_id: String, }, #[serde(untagged, rename_all = "camelCase")] - FormElicitationParamBackwardsCompat { + LegacyForm { #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - meta: Option, + meta: Option, message: String, requested_schema: ElicitationSchema, }, } -impl TryFrom for ElicitRequestParams { +impl TryFrom for ElicitRequestParams { type Error = serde_json::Error; - fn try_from( - value: CreateElicitationRequestParamDeserializeHelper, - ) -> Result { + fn try_from(value: ElicitRequestParamsWire) -> Result { match value { - CreateElicitationRequestParamDeserializeHelper::FormElicitationParam { + ElicitRequestParamsWire::Form { meta, message, requested_schema, } - | CreateElicitationRequestParamDeserializeHelper::FormElicitationParamBackwardsCompat { + | ElicitRequestParamsWire::LegacyForm { meta, message, requested_schema, @@ -2725,7 +3711,7 @@ impl TryFrom for ElicitRequestPa message, requested_schema, }), - CreateElicitationRequestParamDeserializeHelper::UrlElicitationParam { + ElicitRequestParamsWire::Url { meta, message, url, @@ -2771,10 +3757,7 @@ impl TryFrom for ElicitRequestPa /// }; /// ``` #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -#[serde( - tag = "mode", - try_from = "CreateElicitationRequestParamDeserializeHelper" -)] +#[serde(tag = "mode", try_from = "ElicitRequestParamsWire")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub enum ElicitRequestParams { @@ -2782,7 +3765,7 @@ pub enum ElicitRequestParams { FormElicitationParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - meta: Option, + meta: Option, /// Human-readable message explaining what input is needed from the user. /// This should be clear and provide sufficient context for the user to understand /// what information they need to provide. @@ -2797,7 +3780,7 @@ pub enum ElicitRequestParams { UrlElicitationParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - meta: Option, + meta: Option, /// Human-readable message explaining what input is needed from the user. /// This should be clear and provide sufficient context for the user to understand /// what information they need to provide. @@ -2812,13 +3795,13 @@ pub enum ElicitRequestParams { } impl RequestParamsMeta for ElicitRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { match self { ElicitRequestParams::FormElicitationParams { meta, .. } => meta.as_ref(), ElicitRequestParams::UrlElicitationParams { meta, .. } => meta.as_ref(), } } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { match self { ElicitRequestParams::FormElicitationParams { meta, .. } => meta, ElicitRequestParams::UrlElicitationParams { meta, .. } => meta, @@ -2826,13 +3809,6 @@ impl RequestParamsMeta for ElicitRequestParams { } } -/// Deprecated: Use [`ElicitRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use ElicitRequestParams instead")] -pub type CreateElicitationRequestParam = ElicitRequestParams; - -#[deprecated(since = "2.0.0", note = "Renamed to ElicitRequestParams")] -pub type CreateElicitationRequestParams = ElicitRequestParams; - /// The result returned by a client in response to an elicitation request. /// /// Contains the user's decision (accept/decline/cancel) and optionally their input data @@ -2853,7 +3829,7 @@ pub struct ElicitResult { /// Optional protocol-level metadata for this result. #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } impl ElicitResult { @@ -2873,49 +3849,15 @@ impl ElicitResult { } /// Set the metadata on this result. - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } } -#[deprecated(since = "2.0.0", note = "Renamed to ElicitResult")] -pub type CreateElicitationResult = ElicitResult; - /// Request type for creating an elicitation to gather user input pub type ElicitRequest = Request; -#[deprecated(since = "2.0.0", note = "Renamed to ElicitRequest")] -pub type CreateElicitationRequest = ElicitRequest; - -/// Notification parameters for an url elicitation completion notification. -#[derive(Default, Debug, Serialize, Deserialize, Clone, PartialEq)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct ElicitationResponseNotificationParam { - pub elicitation_id: String, - #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, -} - -impl ElicitationResponseNotificationParam { - /// Create a new ElicitationResponseNotificationParam. - pub fn new(elicitation_id: impl Into) -> Self { - Self { - elicitation_id: elicitation_id.into(), - meta: None, - } - } -} - -/// Notification sent when an url elicitation process is completed. -pub type ElicitationCompleteNotification = - Notification; - -#[deprecated(since = "2.0.0", note = "Renamed to ElicitationCompleteNotification")] -pub type ElicitationCompletionNotification = ElicitationCompleteNotification; - // ============================================================================= // TOOL EXECUTION RESULTS // ============================================================================= @@ -2924,11 +3866,23 @@ pub type ElicitationCompletionNotification = ElicitationCompleteNotification; /// /// Contains the content returned by the tool execution and an optional /// flag indicating whether the operation resulted in an error. -#[derive(Default, Debug, Serialize, Clone, PartialEq)] +#[derive(Debug, Serialize, Clone, PartialEq)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub struct CallToolResult { + /// Result type discriminator (SEP-2322). Required by the [spec schema] + /// for servers implementing protocol version `2026-07-28`, but optional + /// here because this type also models results from older protocol + /// versions, which do not carry the field: `None` means absent on the + /// wire, and per the spec "the client MUST treat the absent field as + /// `"complete"`". Constructors default to `Some(ResultType::COMPLETE)`; + /// the server handler clears the field when responding to peers that + /// negotiated an older version. + /// + /// [spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub result_type: Option, /// The content returned by the tool (text, images, etc.) #[serde(default)] pub content: Vec, @@ -2940,7 +3894,7 @@ pub struct CallToolResult { pub is_error: Option, /// Optional protocol-level metadata for this result #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } // Custom Deserialize implementation that: @@ -2948,6 +3902,7 @@ pub struct CallToolResult { // 2. Requires at least one known field to be present, so that `CallToolResult` doesn't // greedily match arbitrary JSON objects when used inside `#[serde(untagged)]` enums // (e.g. `ServerResult`), which would shadow `CustomResult`. +// 3. Rejects non-`complete` result types so other `ServerResult` variants can match. impl<'de> Deserialize<'de> for CallToolResult { fn deserialize(deserializer: D) -> Result where @@ -2956,15 +3911,27 @@ impl<'de> Deserialize<'de> for CallToolResult { #[derive(Deserialize)] #[serde(rename_all = "camelCase")] struct Helper { + #[serde(default)] + result_type: Option, content: Option>, structured_content: Option, is_error: Option, #[serde(rename = "_meta")] - meta: Option, + meta: Option, } let helper = Helper::deserialize(deserializer)?; + if helper + .result_type + .as_ref() + .is_some_and(|result_type| !result_type.is_complete()) + { + return Err(serde::de::Error::custom( + "CallToolResult requires resultType to be \"complete\" when present", + )); + } + if helper.content.is_none() && helper.structured_content.is_none() && helper.is_error.is_none() @@ -2977,6 +3944,7 @@ impl<'de> Deserialize<'de> for CallToolResult { } Ok(CallToolResult { + result_type: helper.result_type, content: helper.content.unwrap_or_default(), structured_content: helper.structured_content, is_error: helper.is_error, @@ -2985,10 +3953,23 @@ impl<'de> Deserialize<'de> for CallToolResult { } } +impl Default for CallToolResult { + fn default() -> Self { + CallToolResult { + result_type: Some(ResultType::COMPLETE), + content: Vec::new(), + structured_content: None, + is_error: None, + meta: None, + } + } +} + impl CallToolResult { /// Create a successful tool result with unstructured content pub fn success(content: Vec) -> Self { CallToolResult { + result_type: Some(ResultType::COMPLETE), content, structured_content: None, is_error: Some(false), @@ -3046,6 +4027,7 @@ impl CallToolResult { /// ``` pub fn error(content: Vec) -> Self { CallToolResult { + result_type: Some(ResultType::COMPLETE), content, structured_content: None, is_error: Some(true), @@ -3068,6 +4050,7 @@ impl CallToolResult { /// ``` pub fn structured(value: Value) -> Self { CallToolResult { + result_type: Some(ResultType::COMPLETE), content: vec![ContentBlock::text(value.to_string())], structured_content: Some(value), is_error: Some(false), @@ -3094,6 +4077,7 @@ impl CallToolResult { /// ``` pub fn structured_error(value: Value) -> Self { CallToolResult { + result_type: Some(ResultType::COMPLETE), content: vec![ContentBlock::text(value.to_string())], structured_content: Some(value), is_error: Some(true), @@ -3102,7 +4086,7 @@ impl CallToolResult { } /// Set the metadata on this result - pub fn with_meta(mut self, meta: Option) -> Self { + pub fn with_meta(mut self, meta: Option) -> Self { self.meta = meta; self } @@ -3151,9 +4135,6 @@ const_string!(CallToolRequestMethod = "tools/call"); /// /// Contains the tool name and optional arguments needed to execute /// the tool operation. -/// -/// This implements `TaskAugmentedRequestParamsMeta` as tool calls can be -/// long-running and may benefit from task-based execution. #[derive(Default, Debug, Serialize, Deserialize, Clone, PartialEq)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] @@ -3161,15 +4142,20 @@ const_string!(CallToolRequestMethod = "tools/call"); pub struct CallToolRequestParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// The name of the tool to call pub name: Cow<'static, str>, /// Arguments to pass to the tool (must match the tool's input schema) #[serde(skip_serializing_if = "Option::is_none")] pub arguments: Option, - /// Task metadata for async task management (SEP-1319) + /// Client responses to server-initiated input requests from a previous + /// [`InputRequiredResult`]. Present only when retrying after an incomplete result. #[serde(skip_serializing_if = "Option::is_none")] - pub task: Option, + pub input_responses: Option, + /// Opaque request state echoed back from a previous [`InputRequiredResult`]. + /// Clients MUST return this value exactly as received. + #[serde(skip_serializing_if = "Option::is_none")] + pub request_state: Option, } impl CallToolRequestParams { @@ -3179,7 +4165,8 @@ impl CallToolRequestParams { meta: None, name: name.into(), arguments: None, - task: None, + input_responses: None, + request_state: None, } } @@ -3189,35 +4176,28 @@ impl CallToolRequestParams { self } - /// Sets the task metadata for this tool call. - pub fn with_task(mut self, task: TaskMetadata) -> Self { - self.task = Some(task); + /// Sets the input responses for an MRTR retry. + pub fn with_input_responses(mut self, input_responses: InputResponses) -> Self { + self.input_responses = Some(input_responses); + self + } + + /// Sets the request state for an MRTR retry. + pub fn with_request_state(mut self, request_state: impl Into) -> Self { + self.request_state = Some(request_state.into()); self } } impl RequestParamsMeta for CallToolRequestParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -impl TaskAugmentedRequestParamsMeta for CallToolRequestParams { - fn task(&self) -> Option<&TaskMetadata> { - self.task.as_ref() - } - fn task_mut(&mut self) -> &mut Option { - &mut self.task - } -} - -/// Deprecated: Use [`CallToolRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use CallToolRequestParams instead")] -pub type CallToolRequestParam = CallToolRequestParams; - /// Request to call a specific tool pub type CallToolRequest = Request; @@ -3281,22 +4261,41 @@ impl CreateMessageResult { } } -#[derive(Default, Debug, Serialize, Deserialize, Clone, PartialEq)] +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub struct GetPromptResult { + /// Result type discriminator (SEP-2322). Required by the [spec schema] + /// for servers implementing protocol version `2026-07-28`, but optional + /// here because this type also models results from older protocol + /// versions, which do not carry the field: `None` means absent on the + /// wire, and per the spec "the client MUST treat the absent field as + /// `"complete"`". Constructors default to `Some(ResultType::COMPLETE)`; + /// the server handler clears the field when responding to peers that + /// negotiated an older version. + /// + /// [spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub result_type: Option, #[serde(skip_serializing_if = "Option::is_none")] pub description: Option, pub messages: Vec, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, +} + +impl Default for GetPromptResult { + fn default() -> Self { + Self::new(Vec::new()) + } } impl GetPromptResult { /// Create a new GetPromptResult with required fields. pub fn new(messages: Vec) -> Self { Self { + result_type: Some(ResultType::COMPLETE), description: None, messages, meta: None, @@ -3311,25 +4310,20 @@ impl GetPromptResult { } // ============================================================================= -// TASK MANAGEMENT +// TASK MANAGEMENT (SEP-2663 Tasks extension: `io.modelcontextprotocol/tasks`) // ============================================================================= const_string!(GetTaskMethod = "tasks/get"); pub type GetTaskRequest = Request; -#[deprecated(since = "2.0.0", note = "Renamed to GetTaskMethod")] -pub type GetTaskInfoMethod = GetTaskMethod; -#[deprecated(since = "2.0.0", note = "Renamed to GetTaskRequest")] -pub type GetTaskInfoRequest = GetTaskRequest; - #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub struct GetTaskParams { - /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, + /// Identifier of the task to query. pub task_id: String, } @@ -3343,65 +4337,53 @@ impl GetTaskParams { } impl RequestParamsMeta for GetTaskParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -#[deprecated(since = "2.0.0", note = "Renamed to GetTaskParams")] -pub type GetTaskInfoParams = GetTaskParams; - -#[deprecated(since = "0.13.0", note = "Use GetTaskParams instead")] -pub type GetTaskInfoParam = GetTaskParams; - -const_string!(ListTasksMethod = "tasks/list"); -pub type ListTasksRequest = RequestOptionalParam; - -const_string!(GetTaskPayloadMethod = "tasks/result"); -pub type GetTaskPayloadRequest = Request; - -#[deprecated(since = "2.0.0", note = "Renamed to GetTaskPayloadMethod")] -pub type GetTaskResultMethod = GetTaskPayloadMethod; -#[deprecated(since = "2.0.0", note = "Renamed to GetTaskPayloadRequest")] -pub type GetTaskResultRequest = GetTaskPayloadRequest; +const_string!(UpdateTaskMethod = "tasks/update"); +pub type UpdateTaskRequest = Request; +/// Parameters for `tasks/update` (SEP-2663): deliver responses to outstanding +/// in-task server-to-client requests surfaced via `tasks/get` `inputRequests`. #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] -pub struct GetTaskPayloadParams { - /// Protocol-level metadata for this request (SEP-1319) +pub struct UpdateTaskParams { #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, + /// Identifier of the task to update. pub task_id: String, + /// Responses to outstanding `inputRequests` previously surfaced by the + /// server. Each key MUST correspond to a currently-outstanding + /// `inputRequests` key. + pub input_responses: InputResponses, } -impl GetTaskPayloadParams { - pub fn new(task_id: impl Into) -> Self { +impl UpdateTaskParams { + pub fn new(task_id: impl Into, input_responses: InputResponses) -> Self { Self { meta: None, task_id: task_id.into(), + input_responses, } } } -impl RequestParamsMeta for GetTaskPayloadParams { - fn meta(&self) -> Option<&Meta> { +impl RequestParamsMeta for UpdateTaskParams { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -#[deprecated(since = "2.0.0", note = "Renamed to GetTaskPayloadParams")] -pub type GetTaskResultParams = GetTaskPayloadParams; -#[deprecated(since = "2.0.0", note = "Renamed to GetTaskPayloadParams")] -pub type GetTaskResultParam = GetTaskPayloadParams; - const_string!(CancelTaskMethod = "tasks/cancel"); pub type CancelTaskRequest = Request; @@ -3412,7 +4394,7 @@ pub type CancelTaskRequest = Request; pub struct CancelTaskParams { /// Protocol-level metadata for this request (SEP-1319) #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, pub task_id: String, } @@ -3426,95 +4408,68 @@ impl CancelTaskParams { } impl RequestParamsMeta for CancelTaskParams { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } -/// Deprecated: Use [`CancelTaskParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use CancelTaskParams instead")] -pub type CancelTaskParam = CancelTaskParams; - // --------------------------------------------------------------------------- -// Task status notification (spec `notifications/tasks/status`) +// Task status notification (SEP-2663 `notifications/tasks`) // --------------------------------------------------------------------------- -const_string!(TaskStatusNotificationMethod = "notifications/tasks/status"); +const_string!(TaskStatusNotificationMethod = "notifications/tasks"); /// Parameters for a task status notification (spec `TaskStatusNotificationParams`). /// -/// The task fields are flattened at the top level: `NotificationParams & Task`. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] +/// Carries a complete [`DetailedTask`] for the current status, identical to +/// what `tasks/get` would have returned at that moment. The task fields are +/// flattened at the top level: `NotificationParams & Task`. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] -pub struct TaskStatusNotificationParam { +pub struct TaskStatusNotificationParams { #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, #[serde(flatten)] - pub task: crate::model::Task, + pub task: crate::model::DetailedTask, } -impl TaskStatusNotificationParam { - pub fn new(task: crate::model::Task) -> Self { +impl TaskStatusNotificationParams { + pub fn new(task: crate::model::DetailedTask) -> Self { Self { meta: None, task } } - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: NotificationMetaObject) -> Self { self.meta = Some(meta); self } } -impl From for TaskStatusNotificationParam { - fn from(task: crate::model::Task) -> Self { +impl From for TaskStatusNotificationParams { + fn from(task: crate::model::DetailedTask) -> Self { Self::new(task) } } -impl Deref for TaskStatusNotificationParam { - type Target = crate::model::Task; +impl Deref for TaskStatusNotificationParams { + type Target = crate::model::DetailedTask; fn deref(&self) -> &Self::Target { &self.task } } -impl DerefMut for TaskStatusNotificationParam { +impl DerefMut for TaskStatusNotificationParams { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.task } } pub type TaskStatusNotification = - Notification; -/// Deprecated: Use [`GetTaskResult`] instead (spec alignment). -#[deprecated(since = "0.15.0", note = "Use GetTaskResult instead")] -pub type GetTaskInfoResult = GetTaskResult; - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct ListTasksResult { - pub tasks: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - pub next_cursor: Option, - #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, -} - -impl ListTasksResult { - pub fn new(tasks: Vec) -> Self { - Self { - tasks, - next_cursor: None, - meta: None, - } - } -} + Notification; // ============================================================================= // MESSAGE TYPE UNIONS @@ -3544,7 +4499,7 @@ macro_rules! ts_union { #[derive(Debug, Serialize, Deserialize, Clone)] #[serde(untagged)] #[allow(clippy::large_enum_variant)] - #[expect(clippy::exhaustive_enums, reason = "intentionally exhaustive")] + #[non_exhaustive] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] pub enum $U { $($declared)* @@ -3574,6 +4529,7 @@ ts_union!( export type ClientRequest = | PingRequest | InitializeRequest + | DiscoverRequest | CompleteRequest | SetLevelRequest | GetPromptRequest @@ -3581,13 +4537,13 @@ ts_union!( | ListResourcesRequest | ListResourceTemplatesRequest | ReadResourceRequest + | SubscriptionsListenRequest | SubscribeRequest | UnsubscribeRequest | CallToolRequest | ListToolsRequest | GetTaskRequest - | ListTasksRequest - | GetTaskPayloadRequest + | UpdateTaskRequest | CancelTaskRequest | CustomRequest; ); @@ -3597,6 +4553,7 @@ impl ClientRequest { match &self { ClientRequest::PingRequest(r) => r.method.as_str(), ClientRequest::InitializeRequest(r) => r.method.as_str(), + ClientRequest::DiscoverRequest(r) => r.method.as_str(), ClientRequest::CompleteRequest(r) => r.method.as_str(), ClientRequest::SetLevelRequest(r) => r.method.as_str(), ClientRequest::GetPromptRequest(r) => r.method.as_str(), @@ -3604,13 +4561,13 @@ impl ClientRequest { ClientRequest::ListResourcesRequest(r) => r.method.as_str(), ClientRequest::ListResourceTemplatesRequest(r) => r.method.as_str(), ClientRequest::ReadResourceRequest(r) => r.method.as_str(), + ClientRequest::SubscriptionsListenRequest(r) => r.method.as_str(), ClientRequest::SubscribeRequest(r) => r.method.as_str(), ClientRequest::UnsubscribeRequest(r) => r.method.as_str(), ClientRequest::CallToolRequest(r) => r.method.as_str(), ClientRequest::ListToolsRequest(r) => r.method.as_str(), ClientRequest::GetTaskRequest(r) => r.method.as_str(), - ClientRequest::ListTasksRequest(r) => r.method.as_str(), - ClientRequest::GetTaskPayloadRequest(r) => r.method.as_str(), + ClientRequest::UpdateTaskRequest(r) => r.method.as_str(), ClientRequest::CancelTaskRequest(r) => r.method.as_str(), ClientRequest::CustomRequest(r) => r.method.as_str(), } @@ -3623,7 +4580,6 @@ ts_union!( | ProgressNotification | InitializedNotification | RootsListChangedNotification - | TaskStatusNotification | CustomNotification; ); @@ -3662,13 +4618,14 @@ ts_union!( | ResourceListChangedNotification | ToolListChangedNotification | PromptListChangedNotification - | ElicitationCompleteNotification + | SubscriptionsAcknowledgedNotification | TaskStatusNotification | CustomNotification; ); ts_union!( export type ServerResult = + | DiscoverResult | InitializeResult | CompleteResult | GetPromptResult @@ -3676,14 +4633,17 @@ ts_union!( | ListResourcesResult | ListResourceTemplatesResult | ReadResourceResult + | SubscriptionsListenResult | ListToolsResult | ElicitResult | CreateTaskResult - | ListTasksResult | GetTaskResult - | CancelTaskResult | CallToolResult - | GetTaskPayloadResult + | InputRequiredResult + // TaskAckResult must come after CallToolResult/InputRequiredResult in this + // untagged union: it only carries `resultType`, so it would otherwise + // shadow any result that includes `resultType: "complete"`. + | TaskAckResult | EmptyResult | CustomResult ; @@ -3693,6 +4653,48 @@ impl ServerResult { pub fn empty(_: ()) -> ServerResult { ServerResult::EmptyResult(EmptyResult {}) } + + /// Empty `tasks/update` / `tasks/cancel` acknowledgement carrying the + /// SEP-2322 `resultType: "complete"` discriminator (SEP-2663). + pub fn task_ack(_: ()) -> ServerResult { + ServerResult::TaskAckResult(TaskAckResult::new()) + } + + /// Strip the SEP-2322 `resultType: "complete"` discriminator so the result + /// keeps the wire shape that predates protocol version `2026-07-28`. + /// + /// The server handler calls this before responding to a peer that + /// negotiated an older protocol version, where the field did not exist and + /// strict peers may reject it. Only the `"complete"` value is stripped: + /// results whose discriminator carries meaning (`"input_required"`, + /// `"task"`) are already gated to `2026-07-28`+ sessions, and custom + /// extension values are preserved. + /// + /// # Examples + /// + /// ``` + /// use rmcp::model::{CallToolResult, ServerResult}; + /// + /// let mut result = ServerResult::CallToolResult(CallToolResult::success(vec![])); + /// result.strip_result_type_for_legacy_peer(); + /// + /// let json = serde_json::to_value(&result).unwrap(); + /// assert!(json.get("resultType").is_none()); + /// ``` + pub fn strip_result_type_for_legacy_peer(&mut self) { + let result_type = match self { + ServerResult::CompleteResult(r) => &mut r.result_type, + ServerResult::GetPromptResult(r) => &mut r.result_type, + ServerResult::ListPromptsResult(r) => &mut r.result_type, + ServerResult::ListResourcesResult(r) => &mut r.result_type, + ServerResult::ListResourceTemplatesResult(r) => &mut r.result_type, + ServerResult::ReadResourceResult(r) => &mut r.result_type, + ServerResult::ListToolsResult(r) => &mut r.result_type, + ServerResult::CallToolResult(r) => &mut r.result_type, + _ => return, + }; + result_type.take_if(|result_type| result_type.is_complete()); + } } pub type ServerJsonRpcMessage = JsonRpcMessage; @@ -3730,12 +4732,61 @@ mod tests { use super::*; #[test] - #[allow(deprecated)] - fn deprecated_aliases_still_resolve() { - // 하위호환: 구 이름이 새 타입으로 여전히 resolve되는지 확인. - let _: CreateElicitationResult = ElicitResult::new(ElicitationAction::Accept); - let _: GetTaskResultParams = GetTaskPayloadParams::new("task-1"); - let _: ResourceReference = ResourceTemplateReference::new("res://x"); + fn known_versions_are_ordered_oldest_first() { + // `known_up_to` walks the list as a sorted prefix. + assert!( + ProtocolVersion::KNOWN_VERSIONS + .windows(2) + .all(|pair| pair[0].as_str() < pair[1].as_str()) + ); + } + + #[test] + fn known_up_to_includes_the_ceiling_itself() { + assert_eq!( + ProtocolVersion::known_up_to(&ProtocolVersion::V_2024_11_05), + &[ProtocolVersion::V_2024_11_05] + ); + } + + #[test] + fn known_up_to_the_newest_version_yields_every_known_version() { + assert_eq!( + ProtocolVersion::known_up_to(&ProtocolVersion::V_2026_07_28), + ProtocolVersion::KNOWN_VERSIONS + ); + } + + #[test] + fn known_up_to_an_unknown_ceiling_stops_at_the_versions_below_it() { + let unknown = ProtocolVersion(Cow::Borrowed("2025-07-01")); + assert_eq!( + ProtocolVersion::known_up_to(&unknown), + &[ + ProtocolVersion::V_2024_11_05, + ProtocolVersion::V_2025_03_26, + ProtocolVersion::V_2025_06_18, + ] + ); + } + + #[test] + fn known_up_to_a_ceiling_below_every_known_version_is_empty() { + let ancient = ProtocolVersion(Cow::Borrowed("1999-01-01")); + assert!(ProtocolVersion::known_up_to(&ancient).is_empty()); + } + + #[cfg(feature = "transport-streamable-http-client")] + #[test] + fn transport_closed_marker_accepts_only_the_process_local_token() { + let local = ErrorData::transport_closed("closed"); + let spoofed = ErrorData::internal_error( + "spoofed", + Some(json!({ "io.modelcontextprotocol/transportClosed": true })), + ); + + assert!(local.is_transport_closed()); + assert!(!spoofed.is_transport_closed()); } #[test] @@ -3945,12 +4996,11 @@ mod tests { serde_json::from_value(request.clone()).expect("invalid request"); let (request, id) = request.into_request().expect("should be a request"); assert_eq!(id, RequestId::Number(1)); - #[allow(deprecated)] match request { ClientRequest::InitializeRequest(Request { method: _, params: - InitializeRequestParam { + InitializeRequestParams { meta: _, protocol_version: _, capabilities, @@ -4216,8 +5266,7 @@ mod tests { } #[test] - fn test_elicitation_deserialization_untagged() { - // Test deserialization without the "type" field (should default to FormElicitationParam) + fn elicitation_without_mode_deserializes_as_form() { let json_data_without_tag = json!({ "message": "Please provide more details.", "requestedSchema": { @@ -4243,7 +5292,7 @@ mod tests { assert_eq!(requested_schema.title, Some(Cow::from("User Details"))); assert_eq!(requested_schema.type_, ObjectTypeConst); } else { - panic!("Expected FormElicitationParam"); + panic!("Expected FormElicitationParams"); } } @@ -4273,13 +5322,15 @@ mod tests { { assert_eq!( meta, - Some(Meta(object!({ "meta_form_key_1": "meta form value 1" }))) + Some(RequestMetaObject(MetaObject( + object!({ "meta_form_key_1": "meta form value 1" }) + ))) ); assert_eq!(message, "Please provide more details."); assert_eq!(requested_schema.title, Some(Cow::from("User Details"))); assert_eq!(requested_schema.type_, ObjectTypeConst); } else { - panic!("Expected FormElicitationParam"); + panic!("Expected FormElicitationParams"); } let json_data_url = json!({ @@ -4300,20 +5351,24 @@ mod tests { { assert_eq!( meta, - Some(Meta(object!({ "meta_url_key_1": "meta url value 1" }))) + Some(RequestMetaObject(MetaObject( + object!({ "meta_url_key_1": "meta url value 1" }) + ))) ); assert_eq!(message, "Please fill out the form at the following URL."); assert_eq!(url, "https://example.com/form"); assert_eq!(elicitation_id, "elicitation-123"); } else { - panic!("Expected UrlElicitationParam"); + panic!("Expected UrlElicitationParams"); } } #[test] fn test_elicitation_serialization() { let form_elicitation = ElicitRequestParams::FormElicitationParams { - meta: Some(Meta(object!({ "meta_form_key_1": "meta form value 1" }))), + meta: Some(RequestMetaObject(MetaObject( + object!({ "meta_form_key_1": "meta form value 1" }), + ))), message: "Please provide more details.".to_string(), requested_schema: ElicitationSchema::builder() .title("User Details") @@ -4337,7 +5392,9 @@ mod tests { assert_eq!(json_form, expected_form_json); let url_elicitation = ElicitRequestParams::UrlElicitationParams { - meta: Some(Meta(object!({ "meta_url_key_1": "meta url value 1" }))), + meta: Some(RequestMetaObject(MetaObject( + object!({ "meta_url_key_1": "meta url value 1" }), + ))), message: "Please fill out the form at the following URL.".to_string(), url: "https://example.com/form".to_string(), elicitation_id: "elicitation-123".to_string(), diff --git a/crates/rmcp/src/model/annotated.rs b/crates/rmcp/src/model/annotated.rs index 06800d1..ac3dbd3 100644 --- a/crates/rmcp/src/model/annotated.rs +++ b/crates/rmcp/src/model/annotated.rs @@ -20,20 +20,23 @@ pub struct Annotations { #[serde(skip_serializing_if = "Option::is_none")] pub priority: Option, #[serde(skip_serializing_if = "Option::is_none", rename = "lastModified")] - pub last_modified: Option>, + pub last_modified: Option, } impl Annotations { - /// Creates a new Annotations instance specifically for resources - /// optional priority, and a timestamp (defaults to now if None) + /// Creates annotations for a resource with priority and an RFC 3339 timestamp. + /// + /// # Panics + /// + /// Panics if `priority` is not in the inclusive range `0.0..=1.0`. pub fn for_resource(priority: f32, timestamp: DateTime) -> Self { assert!( (0.0..=1.0).contains(&priority), "Priority {priority} must be between 0.0 and 1.0" ); - Annotations { + Self { priority: Some(priority), - last_modified: Some(timestamp), + last_modified: Some(timestamp.to_rfc3339()), audience: None, } } @@ -48,12 +51,69 @@ impl Annotations { self } + /// Sets `lastModified` from a typed timestamp, serialized as RFC 3339. pub fn with_timestamp(mut self, timestamp: DateTime) -> Self { - self.last_modified = Some(timestamp); + self.last_modified = Some(timestamp.to_rfc3339()); self } + /// Sets `lastModified` to the current time, serialized as RFC 3339. pub fn with_timestamp_now(self) -> Self { self.with_timestamp(Utc::now()) } } + +#[cfg(test)] +mod tests { + use chrono::TimeZone; + use serde_json::json; + + use super::*; + + fn annotations_of(value: serde_json::Value) -> Annotations { + serde_json::from_value::(value).unwrap() + } + + #[test] + fn preserves_date_only_last_modified_verbatim() { + let annotations = annotations_of(json!({ "lastModified": "2025-01-12" })); + assert_eq!(annotations.last_modified.as_deref(), Some("2025-01-12")); + } + + #[test] + fn preserves_rfc3339_last_modified_verbatim() { + let value = "2025-01-12T15:00:58Z"; + let annotations = annotations_of(json!({ "lastModified": value })); + assert_eq!(annotations.last_modified.as_deref(), Some(value)); + } + + #[test] + fn missing_last_modified_is_none() { + let annotations = annotations_of(json!({})); + assert_eq!(annotations.last_modified, None); + } + + #[test] + fn null_last_modified_is_none() { + let annotations = annotations_of(json!({ "lastModified": null })); + assert_eq!(annotations.last_modified, None); + } + + #[test] + fn invalid_last_modified_string_is_preserved() { + let annotations = annotations_of(json!({ "lastModified": "not-a-date" })); + assert_eq!(annotations.last_modified.as_deref(), Some("not-a-date")); + } + + #[test] + fn timestamp_round_trips_as_rfc3339_string() { + let timestamp = Utc.with_ymd_and_hms(2025, 1, 12, 15, 0, 58).unwrap(); + let annotations = Annotations::default().with_timestamp(timestamp); + + let value = serde_json::to_value(&annotations).unwrap(); + assert_eq!(value["lastModified"], json!(timestamp.to_rfc3339())); + + let round_tripped: Annotations = serde_json::from_value(value).unwrap(); + assert_eq!(round_tripped, annotations); + } +} diff --git a/crates/rmcp/src/model/capabilities.rs b/crates/rmcp/src/model/capabilities.rs index b42e40c..f014f56 100644 --- a/crates/rmcp/src/model/capabilities.rs +++ b/crates/rmcp/src/model/capabilities.rs @@ -72,130 +72,6 @@ pub struct RootsCapabilities { pub list_changed: Option, } -/// Task capabilities shared by client and server. -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct TasksCapability { - #[serde(skip_serializing_if = "Option::is_none")] - pub requests: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub list: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cancel: Option, -} - -/// Request types that support task-augmented execution. -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct TaskRequestsCapability { - #[serde(skip_serializing_if = "Option::is_none")] - pub sampling: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub elicitation: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub tools: Option, -} - -/// Sampling task capability. Deprecated by SEP-2577; remains functional and -/// will be removed in a future release. -/// See . -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct SamplingTaskCapability { - #[serde(skip_serializing_if = "Option::is_none")] - pub create_message: Option, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct ElicitationTaskCapability { - #[serde(skip_serializing_if = "Option::is_none")] - pub create: Option, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct ToolsTaskCapability { - #[serde(skip_serializing_if = "Option::is_none")] - pub call: Option, -} - -impl TasksCapability { - /// Default client tasks capability with sampling and elicitation support. - pub fn client_default() -> Self { - Self { - list: Some(JsonObject::new()), - cancel: Some(JsonObject::new()), - requests: Some(TaskRequestsCapability { - sampling: Some(SamplingTaskCapability { - create_message: Some(JsonObject::new()), - }), - elicitation: Some(ElicitationTaskCapability { - create: Some(JsonObject::new()), - }), - tools: None, - }), - } - } - - /// Default server tasks capability with tools/call support. - pub fn server_default() -> Self { - Self { - list: Some(JsonObject::new()), - cancel: Some(JsonObject::new()), - requests: Some(TaskRequestsCapability { - sampling: None, - elicitation: None, - tools: Some(ToolsTaskCapability { - call: Some(JsonObject::new()), - }), - }), - } - } - - pub fn supports_list(&self) -> bool { - self.list.is_some() - } - - pub fn supports_cancel(&self) -> bool { - self.cancel.is_some() - } - - pub fn supports_tools_call(&self) -> bool { - self.requests - .as_ref() - .and_then(|r| r.tools.as_ref()) - .and_then(|t| t.call.as_ref()) - .is_some() - } - - pub fn supports_sampling_create_message(&self) -> bool { - self.requests - .as_ref() - .and_then(|r| r.sampling.as_ref()) - .and_then(|s| s.create_message.as_ref()) - .is_some() - } - - pub fn supports_elicitation_create(&self) -> bool { - self.requests - .as_ref() - .and_then(|r| r.elicitation.as_ref()) - .and_then(|e| e.create.as_ref()) - .is_some() - } -} - /// Capability for handling elicitation requests from servers. /// Elicitation allows servers to request interactive input from users during tool execution. /// This capability indicates that a client can handle elicitation requests and present @@ -316,8 +192,16 @@ pub struct ClientCapabilities { /// Capability to handle elicitation requests from servers for interactive user input #[serde(skip_serializing_if = "Option::is_none")] pub elicitation: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub tasks: Option, +} + +impl ClientCapabilities { + /// Returns `true` if the `io.modelcontextprotocol/tasks` extension + /// (SEP-2663) is declared in [`Self::extensions`]. + pub fn supports_tasks(&self) -> bool { + self.extensions + .as_ref() + .is_some_and(|e| e.contains_key(super::TASKS_EXTENSION_ID)) + } } /// @@ -356,8 +240,16 @@ pub struct ServerCapabilities { pub resources: Option, #[serde(skip_serializing_if = "Option::is_none")] pub tools: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub tasks: Option, +} + +impl ServerCapabilities { + /// Returns `true` if the `io.modelcontextprotocol/tasks` extension + /// (SEP-2663) is declared in [`Self::extensions`]. + pub fn supports_tasks(&self) -> bool { + self.extensions + .as_ref() + .is_some_and(|e| e.contains_key(super::TASKS_EXTENSION_ID)) + } } #[cfg(any(feature = "server", feature = "macros"))] @@ -484,20 +376,12 @@ builder! { prompts: PromptsCapability, resources: ResourcesCapability, tools: ToolsCapability, - tasks: TasksCapability } } #[cfg(any(feature = "server", feature = "macros"))] -impl< - const E: bool, - const EXT: bool, - const L: bool, - const C: bool, - const P: bool, - const R: bool, - const TASKS: bool, -> ServerCapabilitiesBuilder> +impl + ServerCapabilitiesBuilder> { pub fn enable_tool_list_changed(mut self) -> Self { if let Some(c) = self.tools.as_mut() { @@ -508,15 +392,8 @@ impl< } #[cfg(any(feature = "server", feature = "macros"))] -impl< - const E: bool, - const EXT: bool, - const L: bool, - const C: bool, - const R: bool, - const T: bool, - const TASKS: bool, -> ServerCapabilitiesBuilder> +impl + ServerCapabilitiesBuilder> { pub fn enable_prompts_list_changed(mut self) -> Self { if let Some(c) = self.prompts.as_mut() { @@ -527,15 +404,8 @@ impl< } #[cfg(any(feature = "server", feature = "macros"))] -impl< - const E: bool, - const EXT: bool, - const L: bool, - const C: bool, - const P: bool, - const T: bool, - const TASKS: bool, -> ServerCapabilitiesBuilder> +impl + ServerCapabilitiesBuilder> { pub fn enable_resources_list_changed(mut self) -> Self { if let Some(c) = self.resources.as_mut() { @@ -552,6 +422,18 @@ impl< } } +#[cfg(any(feature = "server", feature = "macros"))] +impl ServerCapabilitiesBuilder { + /// Declare support for the `io.modelcontextprotocol/tasks` extension + /// (SEP-2663) in the `extensions` capability map. + pub fn enable_tasks(mut self) -> Self { + self.extensions + .get_or_insert_with(ExtensionCapabilities::new) + .insert(super::TASKS_EXTENSION_ID.to_string(), JsonObject::new()); + self + } +} + #[cfg(any(feature = "server", feature = "macros"))] builder! { ClientCapabilities{ @@ -568,13 +450,24 @@ builder! { )] sampling: SamplingCapability, elicitation: ElicitationCapability, - tasks: TasksCapability, } } #[cfg(any(feature = "server", feature = "macros"))] -impl - ClientCapabilitiesBuilder> +impl ClientCapabilitiesBuilder { + /// Declare support for the `io.modelcontextprotocol/tasks` extension + /// (SEP-2663) in the `extensions` capability map. + pub fn enable_tasks(mut self) -> Self { + self.extensions + .get_or_insert_with(ExtensionCapabilities::new) + .insert(super::TASKS_EXTENSION_ID.to_string(), JsonObject::new()); + self + } +} + +#[cfg(any(feature = "server", feature = "macros"))] +impl + ClientCapabilitiesBuilder> { #[deprecated( since = "1.8.0", @@ -589,8 +482,8 @@ impl - ClientCapabilitiesBuilder> +impl + ClientCapabilitiesBuilder> { /// Enable tool calling in sampling requests #[deprecated( @@ -618,8 +511,8 @@ impl - ClientCapabilitiesBuilder> +impl + ClientCapabilitiesBuilder> { /// Enable JSON Schema validation for elicitation responses in form mode. /// When enabled, the client will validate user input against the requested_schema @@ -679,68 +572,25 @@ mod test { } #[test] - fn test_task_capabilities_deserialization() { - // Test deserializing from the MCP spec format - let json = serde_json::json!({ - "list": {}, - "cancel": {}, - "requests": { - "tools": { "call": {} } - } - }); - - let tasks: TasksCapability = serde_json::from_value(json).unwrap(); - assert!(tasks.list.is_some()); - assert!(tasks.cancel.is_some()); - assert!(tasks.requests.is_some()); - let requests = tasks.requests.unwrap(); - assert!(requests.tools.is_some()); - assert!(requests.tools.unwrap().call.is_some()); - } - - #[test] - fn test_tasks_capability_client_default() { - let tasks = TasksCapability::client_default(); - - // Verify structure - assert!(tasks.supports_list()); - assert!(tasks.supports_cancel()); - assert!(tasks.supports_sampling_create_message()); - assert!(tasks.supports_elicitation_create()); - assert!(!tasks.supports_tools_call()); - - // Verify serialization matches expected format - let json = serde_json::to_value(&tasks).unwrap(); - assert_eq!(json["list"], serde_json::json!({})); - assert_eq!(json["cancel"], serde_json::json!({})); + fn test_tasks_extension_capability() { + // SEP-2663: tasks are declared via the extensions map. + let capabilities = ClientCapabilities::builder().enable_tasks().build(); + assert!(capabilities.supports_tasks()); + let json = serde_json::to_value(&capabilities).unwrap(); assert_eq!( - json["requests"]["sampling"]["createMessage"], + json["extensions"][crate::model::TASKS_EXTENSION_ID], serde_json::json!({}) ); + + let server = ServerCapabilities::builder().enable_tasks().build(); + assert!(server.supports_tasks()); + let json = serde_json::to_value(&server).unwrap(); assert_eq!( - json["requests"]["elicitation"]["create"], + json["extensions"][crate::model::TASKS_EXTENSION_ID], serde_json::json!({}) ); } - #[test] - fn test_tasks_capability_server_default() { - let tasks = TasksCapability::server_default(); - - // Verify structure - assert!(tasks.supports_list()); - assert!(tasks.supports_cancel()); - assert!(tasks.supports_tools_call()); - assert!(!tasks.supports_sampling_create_message()); - assert!(!tasks.supports_elicitation_create()); - - // Verify serialization matches expected format - let json = serde_json::to_value(&tasks).unwrap(); - assert_eq!(json["list"], serde_json::json!({})); - assert_eq!(json["cancel"], serde_json::json!({})); - assert_eq!(json["requests"]["tools"]["call"], serde_json::json!({})); - } - #[test] #[allow(deprecated)] fn test_client_extensions_capability() { diff --git a/crates/rmcp/src/model/content.rs b/crates/rmcp/src/model/content.rs index 680136f..11b2f0e 100644 --- a/crates/rmcp/src/model/content.rs +++ b/crates/rmcp/src/model/content.rs @@ -10,9 +10,9 @@ // ToolUseContent/ToolResultContent are SEP-2577-deprecated; internal references are expected. #![expect(deprecated)] use serde::{Deserialize, Serialize}; -use serde_json::json; +use serde_json::{Value, json}; -use super::{Annotations, Meta, resource::ResourceContents}; +use super::{Annotations, MetaObject, resource::ResourceContents}; // --------------------------------------------------------------------------- // Flat content structs @@ -28,7 +28,7 @@ pub struct TextContent { pub text: String, /// Optional protocol-level metadata for this content block. #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// Optional annotations describing how the client should use this content. #[serde(skip_serializing_if = "Option::is_none")] pub annotations: Option, @@ -43,7 +43,7 @@ impl TextContent { } } - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } @@ -66,7 +66,7 @@ pub struct ImageContent { pub mime_type: String, /// Optional protocol-level metadata for this content block. #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// Optional annotations describing how the client should use this content. #[serde(skip_serializing_if = "Option::is_none")] pub annotations: Option, @@ -82,7 +82,7 @@ impl ImageContent { } } - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } @@ -105,7 +105,7 @@ pub struct AudioContent { pub mime_type: String, /// Optional protocol-level metadata for this content block. #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// Optional annotations describing how the client should use this content. #[serde(skip_serializing_if = "Option::is_none")] pub annotations: Option, @@ -121,7 +121,7 @@ impl AudioContent { } } - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } @@ -142,7 +142,7 @@ pub struct EmbeddedResource { pub resource: ResourceContents, /// Optional protocol-level metadata for this content block. #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// Optional annotations describing how the client should use this content. #[serde(skip_serializing_if = "Option::is_none")] pub annotations: Option, @@ -164,7 +164,7 @@ impl EmbeddedResource { } } - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } @@ -189,7 +189,7 @@ pub struct ToolUseContent { pub name: String, pub input: super::JsonObject, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } /// Tool execution result in user message (SEP-1577). @@ -203,11 +203,11 @@ pub struct ToolUseContent { )] pub struct ToolResultContent { #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, pub tool_use_id: String, pub content: Vec, #[serde(skip_serializing_if = "Option::is_none")] - pub structured_content: Option, + pub structured_content: Option, #[serde(skip_serializing_if = "Option::is_none")] pub is_error: Option, } @@ -297,7 +297,7 @@ impl ContentBlock { ContentBlock::Resource(EmbeddedResource::new( ResourceContents::TextResourceContents { uri: uri.into(), - mime_type: Some("text".to_string()), + mime_type: Some("text/plain".to_string()), text: content.into(), meta: None, }, @@ -402,7 +402,7 @@ mod tests { #[test] fn test_audio_content_has_meta() { - let audio = AudioContent::new("data", "audio/wav").with_meta(Meta::default()); + let audio = AudioContent::new("data", "audio/wav").with_meta(MetaObject::default()); let json = serde_json::to_value(&audio).unwrap(); assert!(json.get("_meta").is_some()); } diff --git a/crates/rmcp/src/model/elicitation_schema.rs b/crates/rmcp/src/model/elicitation_schema.rs index 3867502..4861087 100644 --- a/crates/rmcp/src/model/elicitation_schema.rs +++ b/crates/rmcp/src/model/elicitation_schema.rs @@ -18,7 +18,8 @@ use std::{borrow::Cow, collections::BTreeMap, marker::PhantomData}; -use serde::{Deserialize, Serialize}; +use indexmap::IndexMap; +use serde::{Deserialize, Deserializer, Serialize}; use crate::{const_string, model::ConstString}; @@ -63,9 +64,6 @@ pub enum PrimitiveSchemaDefinition { Boolean(BooleanSchema), } -#[deprecated(since = "2.0.0", note = "Renamed to PrimitiveSchemaDefinition")] -pub type PrimitiveSchema = PrimitiveSchemaDefinition; - // ============================================================================= // STRING SCHEMA // ============================================================================= @@ -983,17 +981,15 @@ impl EnumSchemaBuilder { return Err("One of the provided default values is not in enum values".to_string()); } } - if let Some(min) = self.min_items { - if (default_values.len() as u64) < min { - return Err("Number of provided default values is less than min_items".to_string()); - } + if let Some(min) = self.min_items + && (default_values.len() as u64) < min + { + return Err("Number of provided default values is less than min_items".to_string()); } - if let Some(max) = self.max_items { - if (default_values.len() as u64) > max { - return Err( - "Number of provided default values is greater than max_items".to_string(), - ); - } + if let Some(max) = self.max_items + && (default_values.len() as u64) > max + { + return Err("Number of provided default values is greater than max_items".to_string()); } self.default = default_values; Ok(self) @@ -1114,11 +1110,21 @@ impl EnumSchema { /// .optional_bool("newsletter", false) /// .build(); /// ``` -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Serialize)] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[serde(rename_all = "camelCase")] +#[cfg_attr(feature = "schemars", schemars(!into))] +#[serde(rename_all = "camelCase", into = "ElicitationSchemaWire")] #[non_exhaustive] pub struct ElicitationSchema { + /// Optional JSON Schema dialect identifier (the `$schema` keyword). + /// + /// The 2025-11-25 protocol revision allows a `requestedSchema` to declare its + /// dialect. It is preserved verbatim so a declared dialect survives a + /// decode/re-encode round-trip instead of being silently dropped. + #[serde(rename = "$schema", default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "schemars", schemars(with = "String"))] + pub schema: Option>, + /// Always "object" for elicitation schemas #[serde(rename = "type")] pub type_: ObjectTypeConst, @@ -1130,6 +1136,11 @@ pub struct ElicitationSchema { /// Property definitions (must be primitive types) pub properties: BTreeMap, + /// Property names in wire order. Schemas constructed from a `BTreeMap` + /// use the map's sorted key order. + #[serde(skip)] + pub property_order: Option>, + /// List of required property names #[serde(skip_serializing_if = "Option::is_none")] pub required: Option>, @@ -1139,13 +1150,80 @@ pub struct ElicitationSchema { pub description: Option>, } +#[derive(Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +struct ElicitationSchemaWire { + #[serde(rename = "$schema", default, skip_serializing_if = "Option::is_none")] + schema: Option>, + #[serde(rename = "type")] + type_: ObjectTypeConst, + #[serde(skip_serializing_if = "Option::is_none")] + title: Option>, + properties: IndexMap, + #[serde(skip_serializing_if = "Option::is_none")] + required: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + description: Option>, +} + +impl From for ElicitationSchema { + fn from(schema: ElicitationSchemaWire) -> Self { + Self { + schema: schema.schema, + type_: schema.type_, + title: schema.title, + property_order: Some(schema.properties.keys().cloned().collect()), + properties: schema.properties.into_iter().collect(), + required: schema.required, + description: schema.description, + } + } +} + +impl From for ElicitationSchemaWire { + fn from(schema: ElicitationSchema) -> Self { + let mut remaining = schema.properties; + let mut properties = IndexMap::with_capacity(remaining.len()); + + if let Some(property_order) = schema.property_order { + for name in property_order { + if let Some(definition) = remaining.remove(&name) { + properties.insert(name, definition); + } + } + } + properties.extend(remaining); + + Self { + schema: schema.schema, + type_: schema.type_, + title: schema.title, + properties, + required: schema.required, + description: schema.description, + } + } +} + +impl<'de> Deserialize<'de> for ElicitationSchema { + fn deserialize<__D>(__deserializer: __D) -> Result + where + __D: Deserializer<'de>, + { + ElicitationSchemaWire::deserialize(__deserializer).map(Into::into) + } +} + impl ElicitationSchema { /// Create a new elicitation schema with the given properties pub fn new(properties: BTreeMap) -> Self { + let property_order = Some(properties.keys().cloned().collect()); Self { + schema: None, type_: ObjectTypeConst, title: None, properties, + property_order, required: None, description: None, } @@ -1237,6 +1315,12 @@ impl ElicitationSchema { self } + /// Set the JSON Schema dialect identifier (the `$schema` keyword). + pub fn with_schema(mut self, schema: impl Into>) -> Self { + self.schema = Some(schema.into()); + self + } + /// Create a builder for constructing elicitation schemas fluently pub fn builder() -> ElicitationSchemaBuilder { ElicitationSchemaBuilder::new() @@ -1602,44 +1686,6 @@ impl ElicitationSchemaBuilder { self.property(name, PrimitiveSchemaDefinition::Enum(enum_schema)) } - /// Add a required enum property using values. Creates an untitled single-select enum. - #[deprecated( - since = "0.13.0", - note = "Use ElicitationSchemaBuilder::required_enum_schema with EnumSchema::builder instead" - )] - pub fn required_enum(self, name: impl Into, values: Vec) -> Self { - self.required_property( - name, - PrimitiveSchemaDefinition::Enum(EnumSchema::Legacy(LegacyEnumSchema { - type_: StringTypeConst, - title: None, - description: None, - enum_: values, - enum_names: None, - default: None, - })), - ) - } - - /// Add an optional enum property using values. Creates an untitled single-select enum. - #[deprecated( - since = "0.13.0", - note = "Use ElicitationSchemaBuilder::optional_enum_schema with EnumSchema::builder instead" - )] - pub fn optional_enum(self, name: impl Into, values: Vec) -> Self { - self.property( - name, - PrimitiveSchemaDefinition::Enum(EnumSchema::Legacy(LegacyEnumSchema { - type_: StringTypeConst, - title: None, - description: None, - enum_: values, - enum_names: None, - default: None, - })), - ) - } - /// Mark an existing property as required pub fn mark_required(mut self, name: impl Into) -> Self { self.required.push(name.into()); @@ -1675,10 +1721,13 @@ impl ElicitationSchemaBuilder { } } + let property_order = Some(self.properties.keys().cloned().collect()); Ok(ElicitationSchema { + schema: None, type_: ObjectTypeConst, title: self.title, properties: self.properties, + property_order, required: if self.required.is_empty() { None } else { @@ -1864,6 +1913,61 @@ mod tests { output["properties"]["choice"]["enumNames"], serde_json::json!(["Option One", "Option Two", "Option Three"]), ); + let input = r#"{"type":"object","properties":{"firstName":{"type":"string"},"lastName":{"type":"string"},"email":{"type":"string"}}}"#; + let ordered: ElicitationSchema = serde_json::from_str(input)?; + assert_eq!( + ordered.property_order.as_ref().unwrap().join(","), + "firstName,lastName,email", + ); + assert_eq!(serde_json::to_string(&ordered)?, input); + Ok(()) + } + + #[test] + fn test_elicitation_schema_preserves_schema_dialect_roundtrip() -> anyhow::Result<()> { + // Regression test for #1168: a top-level `$schema` dialect declaration on a + // `requestedSchema` was silently dropped, because the wire bridge struct had + // no field to hold it. It must survive a decode/re-encode round-trip. + let input = serde_json::json!({ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "name": { "type": "string" } + } + }); + let schema: ElicitationSchema = serde_json::from_value(input.clone())?; + assert_eq!( + schema.schema.as_deref(), + Some("https://json-schema.org/draft/2020-12/schema"), + ); + let output = serde_json::to_value(&schema)?; + assert_eq!(output, input); + Ok(()) + } + + #[test] + fn test_elicitation_schema_omits_schema_dialect_when_absent() -> anyhow::Result<()> { + // A schema with no dialect must not emit a `$schema` key (no `"$schema": null`). + let input = serde_json::json!({ + "type": "object", + "properties": { "name": { "type": "string" } } + }); + let schema: ElicitationSchema = serde_json::from_value(input)?; + assert!(schema.schema.is_none()); + let json = serde_json::to_value(&schema)?; + assert!(json.get("$schema").is_none()); + Ok(()) + } + + #[test] + fn test_elicitation_schema_with_schema_setter_serializes_dialect() -> anyhow::Result<()> { + let schema = ElicitationSchema::new(BTreeMap::new()) + .with_schema("https://json-schema.org/draft/2020-12/schema"); + let json = serde_json::to_value(&schema)?; + assert_eq!( + json["$schema"], + "https://json-schema.org/draft/2020-12/schema" + ); Ok(()) } diff --git a/crates/rmcp/src/model/meta.rs b/crates/rmcp/src/model/meta.rs index 712a439..0a7121e 100644 --- a/crates/rmcp/src/model/meta.rs +++ b/crates/rmcp/src/model/meta.rs @@ -5,13 +5,29 @@ use serde_json::Value; use super::{ ClientCapabilities, ClientNotification, ClientRequest, CustomNotification, CustomRequest, - Extensions, Implementation, JsonObject, JsonRpcMessage, LoggingLevel, NumberOrString, - ProgressToken, ProtocolVersion, ServerNotification, ServerRequest, TaskMetadata, + Extensions, Implementation, JsonObject, JsonRpcMessage, LoggingLevel, ProgressToken, + ProtocolVersion, RequestId, ServerNotification, ServerRequest, }; +/// Access to the metadata carried by a message envelope's [`Extensions`]. +/// +/// The metadata type differs by message kind: requests carry a +/// [`RequestMetaObject`] and notifications carry a [`NotificationMetaObject`]. +/// +/// The envelope extensions are the canonical runtime location for `_meta`: +/// deserialization strips the wire `params._meta` into the extensions (typed +/// params `meta` fields stay empty), and the service loop moves it into +/// [`RequestContext::meta`] / [`NotificationContext::meta`] before dispatch. +/// Typed params `meta` fields are honored when serializing outgoing messages; +/// on key conflicts the extensions-level metadata wins. +/// +/// [`RequestContext::meta`]: crate::service::RequestContext +/// [`NotificationContext::meta`]: crate::service::NotificationContext pub trait GetMeta { - fn get_meta_mut(&mut self) -> &mut Meta; - fn get_meta(&self) -> &Meta; + /// The metadata type for this message kind. + type Metadata: Default; + fn get_meta_mut(&mut self) -> &mut Self::Metadata; + fn get_meta(&self) -> &Self::Metadata; } pub trait GetExtensions { @@ -21,15 +37,16 @@ pub trait GetExtensions { /// Trait for request params that contain the `_meta` field. /// -/// Per the MCP 2025-11-25 spec, all request params should have an optional `_meta` -/// field that can contain a `progressToken` for tracking long-running operations. +/// Per the MCP spec, all request params may have an optional `_meta` +/// field ([`RequestMetaObject`]) that can contain a `progressToken` for +/// tracking long-running operations. pub trait RequestParamsMeta { /// Get a reference to the meta field - fn meta(&self) -> Option<&Meta>; + fn meta(&self) -> Option<&RequestMetaObject>; /// Get a mutable reference to the meta field - fn meta_mut(&mut self) -> &mut Option; + fn meta_mut(&mut self) -> &mut Option; /// Set the meta field - fn set_meta(&mut self, meta: Meta) { + fn set_meta(&mut self, meta: RequestMetaObject) { *self.meta_mut() = Some(meta); } /// Get the progress token from meta, if present @@ -41,7 +58,7 @@ pub trait RequestParamsMeta { match self.meta_mut() { Some(meta) => meta.set_progress_token(token), none => { - let mut meta = Meta::new(); + let mut meta = RequestMetaObject::new(); meta.set_progress_token(token); *none = Some(meta); } @@ -72,23 +89,8 @@ pub trait RequestParamsMeta { self.meta_or_default().set_baggage(value); } /// Get a mutable reference to meta, inserting an empty one if absent. - fn meta_or_default(&mut self) -> &mut Meta { - self.meta_mut().get_or_insert_with(Meta::new) - } -} - -/// Trait for task-augmented request params that contain both `_meta` and `task` fields. -/// -/// Per the MCP 2025-11-25 spec, certain requests (like `tools/call` and `sampling/createMessage`) -/// can include a `task` field to signal that the caller wants task-augmented execution. -pub trait TaskAugmentedRequestParamsMeta: RequestParamsMeta { - /// Get a reference to the task field - fn task(&self) -> Option<&TaskMetadata>; - /// Get a mutable reference to the task field - fn task_mut(&mut self) -> &mut Option; - /// Set the task field - fn set_task(&mut self, task: TaskMetadata) { - *self.task_mut() = Some(task); + fn meta_or_default(&mut self) -> &mut RequestMetaObject { + self.meta_mut().get_or_insert_with(RequestMetaObject::new) } } @@ -102,13 +104,14 @@ impl GetExtensions for CustomNotification { } impl GetMeta for CustomNotification { - fn get_meta_mut(&mut self) -> &mut Meta { + type Metadata = NotificationMetaObject; + fn get_meta_mut(&mut self) -> &mut NotificationMetaObject { self.extensions_mut().get_or_insert_default() } - fn get_meta(&self) -> &Meta { + fn get_meta(&self) -> &NotificationMetaObject { self.extensions() - .get::() - .unwrap_or(Meta::static_empty()) + .get::() + .unwrap_or(NotificationMetaObject::static_empty()) } } @@ -122,19 +125,20 @@ impl GetExtensions for CustomRequest { } impl GetMeta for CustomRequest { - fn get_meta_mut(&mut self) -> &mut Meta { + type Metadata = RequestMetaObject; + fn get_meta_mut(&mut self) -> &mut RequestMetaObject { self.extensions_mut().get_or_insert_default() } - fn get_meta(&self) -> &Meta { + fn get_meta(&self) -> &RequestMetaObject { self.extensions() - .get::() - .unwrap_or(Meta::static_empty()) + .get::() + .unwrap_or(RequestMetaObject::static_empty()) } } macro_rules! variant_extension { ( - $Enum: ident { + $Enum: ident: $Metadata: ident { $($variant: ident)* } ) => { @@ -155,20 +159,22 @@ macro_rules! variant_extension { } } impl GetMeta for $Enum { - fn get_meta_mut(&mut self) -> &mut Meta { + type Metadata = $Metadata; + fn get_meta_mut(&mut self) -> &mut $Metadata { self.extensions_mut().get_or_insert_default() } - fn get_meta(&self) -> &Meta { - self.extensions().get::().unwrap_or(Meta::static_empty()) + fn get_meta(&self) -> &$Metadata { + self.extensions().get::<$Metadata>().unwrap_or($Metadata::static_empty()) } } }; } variant_extension! { - ClientRequest { + ClientRequest: RequestMetaObject { PingRequest InitializeRequest + DiscoverRequest CompleteRequest SetLevelRequest GetPromptRequest @@ -176,20 +182,20 @@ variant_extension! { ListResourcesRequest ListResourceTemplatesRequest ReadResourceRequest + SubscriptionsListenRequest SubscribeRequest UnsubscribeRequest CallToolRequest ListToolsRequest CustomRequest GetTaskRequest - ListTasksRequest - GetTaskPayloadRequest + UpdateTaskRequest CancelTaskRequest } } variant_extension! { - ServerRequest { + ServerRequest: RequestMetaObject { PingRequest CreateMessageRequest ListRootsRequest @@ -199,18 +205,17 @@ variant_extension! { } variant_extension! { - ClientNotification { + ClientNotification: NotificationMetaObject { CancelledNotification ProgressNotification InitializedNotification RootsListChangedNotification - TaskStatusNotification CustomNotification } } variant_extension! { - ServerNotification { + ServerNotification: NotificationMetaObject { CancelledNotification ProgressNotification LoggingMessageNotification @@ -218,23 +223,27 @@ variant_extension! { ResourceListChangedNotification ToolListChangedNotification PromptListChangedNotification - ElicitationCompleteNotification + SubscriptionsAcknowledgedNotification TaskStatusNotification CustomNotification } } + +/// General-purpose `_meta` map (spec `MetaObject`). +/// +/// This is the metadata shape used by results, content blocks, and catalog +/// descriptors (tools, prompts, resources, roots, ...). It preserves arbitrary +/// extension keys and offers helpers for the reserved W3C Trace Context keys +/// (SEP-414). +/// +/// Request and notification `_meta` maps have additional reserved keys; see +/// [`RequestMetaObject`] and [`NotificationMetaObject`]. #[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[serde(transparent)] #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] -pub struct Meta(pub JsonObject); +pub struct MetaObject(pub JsonObject); -impl Meta { - const PROGRESS_TOKEN_FIELD: &str = "progressToken"; - const META_KEY_PROTOCOL_VERSION: &str = "io.modelcontextprotocol/protocolVersion"; - const META_KEY_CLIENT_INFO: &str = "io.modelcontextprotocol/clientInfo"; - const META_KEY_CLIENT_CAPABILITIES: &str = "io.modelcontextprotocol/clientCapabilities"; - const META_KEY_LOG_LEVEL: &str = "io.modelcontextprotocol/logLevel"; +impl MetaObject { /// Reserved `_meta` key for the W3C Trace Context `traceparent` value (SEP-414). const TRACEPARENT_FIELD: &str = "traceparent"; /// Reserved `_meta` key for the W3C Trace Context `tracestate` value (SEP-414). @@ -242,102 +251,11 @@ impl Meta { /// Reserved `_meta` key for the W3C Baggage value (SEP-414). const BAGGAGE_FIELD: &str = "baggage"; + /// Create an empty metadata map. pub fn new() -> Self { Self(JsonObject::new()) } - /// Create a new Meta with a progress token set - pub fn with_progress_token(token: ProgressToken) -> Self { - let mut meta = Self::new(); - meta.set_progress_token(token); - meta - } - - pub(crate) fn static_empty() -> &'static Self { - static EMPTY: std::sync::OnceLock = std::sync::OnceLock::new(); - EMPTY.get_or_init(Default::default) - } - - pub fn get_progress_token(&self) -> Option { - self.0 - .get(Self::PROGRESS_TOKEN_FIELD) - .and_then(|v| match v { - Value::String(s) => { - Some(ProgressToken(NumberOrString::String(s.to_string().into()))) - } - Value::Number(n) => { - if let Some(i) = n.as_i64() { - Some(ProgressToken(NumberOrString::Number(i))) - } else if let Some(u) = n.as_u64() { - if u <= i64::MAX as u64 { - Some(ProgressToken(NumberOrString::Number(u as i64))) - } else { - None - } - } else { - None - } - } - _ => None, - }) - } - - pub fn set_progress_token(&mut self, token: ProgressToken) { - match token.0 { - NumberOrString::String(ref s) => self.0.insert( - Self::PROGRESS_TOKEN_FIELD.to_string(), - Value::String(s.to_string()), - ), - NumberOrString::Number(n) => self.0.insert( - Self::PROGRESS_TOKEN_FIELD.to_string(), - Value::Number(n.into()), - ), - }; - } - - /// Get the MCP protocol version carried in `_meta`, if present and valid. - pub fn protocol_version(&self) -> Option { - self.decode_value(Self::META_KEY_PROTOCOL_VERSION) - } - - /// Set the MCP protocol version carried in `_meta`. - pub fn set_protocol_version(&mut self, protocol_version: ProtocolVersion) { - self.0.insert( - Self::META_KEY_PROTOCOL_VERSION.to_string(), - Value::String(protocol_version.to_string()), - ); - } - - /// Get the client implementation identity carried in `_meta`, if present and valid. - pub fn client_info(&self) -> Option { - self.decode_value(Self::META_KEY_CLIENT_INFO) - } - - /// Set the client implementation identity carried in `_meta`. - pub fn set_client_info(&mut self, client_info: Implementation) { - self.insert_serialized(Self::META_KEY_CLIENT_INFO, client_info); - } - - /// Get the client capabilities carried in `_meta`, if present and valid. - pub fn client_capabilities(&self) -> Option { - self.decode_value(Self::META_KEY_CLIENT_CAPABILITIES) - } - - /// Set the client capabilities carried in `_meta`. - pub fn set_client_capabilities(&mut self, client_capabilities: ClientCapabilities) { - self.insert_serialized(Self::META_KEY_CLIENT_CAPABILITIES, client_capabilities); - } - - /// Get the requested per-request log level carried in `_meta`, if present and valid. - pub fn log_level(&self) -> Option { - self.decode_value(Self::META_KEY_LOG_LEVEL) - } - - /// Set the requested per-request log level carried in `_meta`. - pub fn set_log_level(&mut self, log_level: LoggingLevel) { - self.insert_serialized(Self::META_KEY_LOG_LEVEL, log_level); - } - /// Read a string-valued `_meta` field, or `None` if absent or not a string. fn get_str(&self, field: &str) -> Option<&str> { self.0.get(field).and_then(Value::as_str) @@ -357,9 +275,9 @@ impl Meta { /// Set the W3C `traceparent` value (SEP-414). /// /// ``` - /// use rmcp::model::Meta; + /// use rmcp::model::MetaObject; /// - /// let mut meta = Meta::new(); + /// let mut meta = MetaObject::new(); /// meta.set_traceparent("00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"); /// assert_eq!( /// meta.get_traceparent(), @@ -390,10 +308,9 @@ impl Meta { self.set_str(Self::BAGGAGE_FIELD, value); } - pub fn extend(&mut self, other: Meta) { - for (k, v) in other.0.into_iter() { - self.0.insert(k, v); - } + /// Insert every entry of `other`, overwriting existing keys on conflict. + pub fn extend(&mut self, other: MetaObject) { + self.0.extend(other.0); } fn decode_value(&self, key: &str) -> Option @@ -413,7 +330,7 @@ impl Meta { } } -impl Deref for Meta { +impl Deref for MetaObject { type Target = JsonObject; fn deref(&self) -> &Self::Target { @@ -421,12 +338,362 @@ impl Deref for Meta { } } -impl DerefMut for Meta { +impl DerefMut for MetaObject { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } } +impl From for MetaObject { + fn from(object: JsonObject) -> Self { + Self(object) + } +} + +#[cfg(feature = "schemars")] +impl schemars::JsonSchema for MetaObject { + fn schema_name() -> std::borrow::Cow<'static, str> { + std::borrow::Cow::Borrowed("MetaObject") + } + + fn json_schema(_: &mut schemars::SchemaGenerator) -> schemars::Schema { + schemars::json_schema!({ + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", + "type": "object", + "additionalProperties": true, + }) + } +} + +/// The `_meta` map carried by requests (spec `RequestMetaObject`). +/// +/// In addition to arbitrary extension keys, requests reserve: +/// - `progressToken` for progress tracking +/// - `io.modelcontextprotocol/protocolVersion` (SEP-2575) +/// - `io.modelcontextprotocol/clientInfo` (SEP-2575) +/// - `io.modelcontextprotocol/clientCapabilities` (SEP-2575) +/// - `io.modelcontextprotocol/logLevel` (SEP-2575) +/// +/// The 2026-07-28 draft schema requires the protocol-version and +/// client-capabilities keys; client-info is optional. Earlier protocol versions +/// do not know them. All keys therefore stay optional at runtime and in the +/// generated (version-shared) JSON schema — use +/// [`RequestMetaObject::missing_required_keys`] to validate a request against +/// the negotiated protocol version. +/// +/// This type dereferences to [`MetaObject`] (and transitively to the underlying +/// map), so general helpers such as the SEP-414 trace-context accessors remain +/// available. +#[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)] +#[serde(transparent)] +#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] +pub struct RequestMetaObject(pub MetaObject); + +impl RequestMetaObject { + const PROGRESS_TOKEN_FIELD: &str = "progressToken"; + const META_KEY_PROTOCOL_VERSION: &str = "io.modelcontextprotocol/protocolVersion"; + const META_KEY_CLIENT_INFO: &str = "io.modelcontextprotocol/clientInfo"; + const META_KEY_CLIENT_CAPABILITIES: &str = "io.modelcontextprotocol/clientCapabilities"; + const META_KEY_LOG_LEVEL: &str = "io.modelcontextprotocol/logLevel"; + + /// Request `_meta` keys the 2026-07-28 draft schema marks as required. + pub const DRAFT_REQUIRED_KEYS: [&str; 2] = [ + Self::META_KEY_PROTOCOL_VERSION, + Self::META_KEY_CLIENT_CAPABILITIES, + ]; + + /// Create an empty request metadata map. + pub fn new() -> Self { + Self::default() + } + + /// Create a new request meta with a progress token set + pub fn with_progress_token(token: ProgressToken) -> Self { + let mut meta = Self::new(); + meta.set_progress_token(token); + meta + } + + /// Create request metadata with the client context SEP-2575 requires on every request. + pub fn with_client_context( + protocol_version: ProtocolVersion, + client_info: Implementation, + client_capabilities: ClientCapabilities, + ) -> Self { + let mut meta = Self::new(); + meta.set_protocol_version(protocol_version); + meta.set_client_info(client_info); + meta.set_client_capabilities(client_capabilities); + meta + } + + pub(crate) fn static_empty() -> &'static Self { + static EMPTY: std::sync::OnceLock = std::sync::OnceLock::new(); + EMPTY.get_or_init(Default::default) + } + + /// Get the progress token carried in `_meta`, if present and valid. + pub fn get_progress_token(&self) -> Option { + self.0.decode_value(Self::PROGRESS_TOKEN_FIELD) + } + + /// Set the progress token carried in `_meta`. + pub fn set_progress_token(&mut self, token: ProgressToken) { + self.0.insert_serialized(Self::PROGRESS_TOKEN_FIELD, token); + } + + /// Get the MCP protocol version carried in `_meta`, if present and valid. + pub fn protocol_version(&self) -> Option { + self.0.decode_value(Self::META_KEY_PROTOCOL_VERSION) + } + + /// Set the MCP protocol version carried in `_meta`. + pub fn set_protocol_version(&mut self, protocol_version: ProtocolVersion) { + self.0.0.insert( + Self::META_KEY_PROTOCOL_VERSION.to_string(), + Value::String(protocol_version.to_string()), + ); + } + + /// Get the client implementation identity carried in `_meta`, if present and valid. + pub fn client_info(&self) -> Option { + self.0.decode_value(Self::META_KEY_CLIENT_INFO) + } + + /// Set the client implementation identity carried in `_meta`. + pub fn set_client_info(&mut self, client_info: Implementation) { + self.0 + .insert_serialized(Self::META_KEY_CLIENT_INFO, client_info); + } + + /// Get the client capabilities carried in `_meta`, if present and valid. + pub fn client_capabilities(&self) -> Option { + self.0.decode_value(Self::META_KEY_CLIENT_CAPABILITIES) + } + + /// Set the client capabilities carried in `_meta`. + pub fn set_client_capabilities(&mut self, client_capabilities: ClientCapabilities) { + self.0 + .insert_serialized(Self::META_KEY_CLIENT_CAPABILITIES, client_capabilities); + } + + /// Get the requested per-request log level carried in `_meta`, if present and valid. + pub fn log_level(&self) -> Option { + self.0.decode_value(Self::META_KEY_LOG_LEVEL) + } + + /// Set the requested per-request log level carried in `_meta`. + pub fn set_log_level(&mut self, log_level: LoggingLevel) { + self.0 + .insert_serialized(Self::META_KEY_LOG_LEVEL, log_level); + } + + /// Return the [`Self::DRAFT_REQUIRED_KEYS`] whose values are absent or + /// invalid in this map, if `protocol_version` requires them. + /// + /// A key counts as missing when it is not present *or* when its value does + /// not decode into the expected type (e.g. a numeric `protocolVersion` or + /// a string `clientInfo`), matching what the typed accessors return. + /// + /// Protocol versions before 2026-07-28 have no required request metadata, + /// so this always returns an empty list for them. + /// + /// # Examples + /// + /// ``` + /// use rmcp::model::{ProtocolVersion, RequestMetaObject}; + /// + /// let meta = RequestMetaObject::new(); + /// // Older protocols have no required request metadata. + /// assert!( + /// meta.missing_required_keys(&ProtocolVersion::V_2025_11_25) + /// .is_empty() + /// ); + /// // The 2026-07-28 protocol requires per-request context. + /// assert_eq!( + /// meta.missing_required_keys(&ProtocolVersion::V_2026_07_28), + /// RequestMetaObject::DRAFT_REQUIRED_KEYS.to_vec(), + /// ); + /// ``` + pub fn missing_required_keys(&self, protocol_version: &ProtocolVersion) -> Vec<&'static str> { + if protocol_version.as_str() < ProtocolVersion::V_2026_07_28.as_str() { + return Vec::new(); + } + let mut missing = Vec::new(); + if self.protocol_version().is_none() { + missing.push(Self::META_KEY_PROTOCOL_VERSION); + } + if self.client_capabilities().is_none() { + missing.push(Self::META_KEY_CLIENT_CAPABILITIES); + } + missing + } + + /// Insert every entry of `other`, overwriting existing keys on conflict. + pub fn extend(&mut self, other: RequestMetaObject) { + self.0.extend(other.0); + } +} + +impl Deref for RequestMetaObject { + type Target = MetaObject; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for RequestMetaObject { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl From for RequestMetaObject { + fn from(meta: MetaObject) -> Self { + Self(meta) + } +} + +impl From for RequestMetaObject { + fn from(object: JsonObject) -> Self { + Self(MetaObject(object)) + } +} + +#[cfg(feature = "schemars")] +impl schemars::JsonSchema for RequestMetaObject { + fn schema_name() -> std::borrow::Cow<'static, str> { + std::borrow::Cow::Borrowed("RequestMetaObject") + } + + fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema { + let progress_token = generator.subschema_for::(); + let client_info = generator.subschema_for::(); + let client_capabilities = generator.subschema_for::(); + let log_level = generator.subschema_for::(); + // rmcp generates one schema shared by every supported protocol + // version, so the keys validated for 2026-07-28 are left + // optional here: a 2025-11-25 request whose `_meta` only carries + // `progressToken` is valid. Version-specific validation is available at + // runtime via [`RequestMetaObject::missing_required_keys`]. + schemars::json_schema!({ + "description": "Metadata reserved by MCP on requests. Extension keys are also allowed.", + "type": "object", + "properties": { + "progressToken": progress_token, + "io.modelcontextprotocol/protocolVersion": { + "type": "string", + }, + "io.modelcontextprotocol/clientInfo": client_info, + "io.modelcontextprotocol/clientCapabilities": client_capabilities, + "io.modelcontextprotocol/logLevel": log_level, + }, + "additionalProperties": true, + }) + } +} + +/// The `_meta` map carried by notifications (spec `NotificationMetaObject`). +/// +/// In addition to arbitrary extension keys, notifications reserve +/// `io.modelcontextprotocol/subscriptionId` to correlate a notification with a +/// prior subscription request. +/// +/// This type dereferences to [`MetaObject`] (and transitively to the underlying +/// map), so general helpers such as the SEP-414 trace-context accessors remain +/// available. +#[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)] +#[serde(transparent)] +#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] +pub struct NotificationMetaObject(pub MetaObject); + +impl NotificationMetaObject { + const META_KEY_SUBSCRIPTION_ID: &str = "io.modelcontextprotocol/subscriptionId"; + + /// Create an empty notification metadata map. + pub fn new() -> Self { + Self::default() + } + + pub(crate) fn static_empty() -> &'static Self { + static EMPTY: std::sync::OnceLock = std::sync::OnceLock::new(); + EMPTY.get_or_init(Default::default) + } + + /// Get the subscription id carried in `_meta`, if present and valid. + /// + /// # Examples + /// + /// ``` + /// use rmcp::model::{NotificationMetaObject, RequestId}; + /// + /// let mut meta = NotificationMetaObject::new(); + /// assert_eq!(meta.subscription_id(), None); + /// meta.set_subscription_id(RequestId::Number(7)); + /// assert_eq!(meta.subscription_id(), Some(RequestId::Number(7))); + /// ``` + pub fn subscription_id(&self) -> Option { + self.0.decode_value(Self::META_KEY_SUBSCRIPTION_ID) + } + + /// Set the subscription id carried in `_meta`. + pub fn set_subscription_id(&mut self, subscription_id: RequestId) { + self.0 + .insert_serialized(Self::META_KEY_SUBSCRIPTION_ID, subscription_id); + } + + /// Insert every entry of `other`, overwriting existing keys on conflict. + pub fn extend(&mut self, other: NotificationMetaObject) { + self.0.extend(other.0); + } +} + +impl Deref for NotificationMetaObject { + type Target = MetaObject; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for NotificationMetaObject { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl From for NotificationMetaObject { + fn from(meta: MetaObject) -> Self { + Self(meta) + } +} + +impl From for NotificationMetaObject { + fn from(object: JsonObject) -> Self { + Self(MetaObject(object)) + } +} + +#[cfg(feature = "schemars")] +impl schemars::JsonSchema for NotificationMetaObject { + fn schema_name() -> std::borrow::Cow<'static, str> { + std::borrow::Cow::Borrowed("NotificationMetaObject") + } + + fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema { + let subscription_id = generator.subschema_for::(); + schemars::json_schema!({ + "description": "Metadata reserved by MCP on notifications. Extension keys are also allowed.", + "type": "object", + "properties": { + "io.modelcontextprotocol/subscriptionId": subscription_id, + }, + "additionalProperties": true, + }) + } +} + impl JsonRpcMessage where Req: GetExtensions, @@ -451,17 +718,18 @@ where #[cfg(test)] mod tests { use super::*; + use crate::model::NumberOrString; #[derive(Default)] struct Params { - meta: Option, + meta: Option, } impl RequestParamsMeta for Params { - fn meta(&self) -> Option<&Meta> { + fn meta(&self) -> Option<&RequestMetaObject> { self.meta.as_ref() } - fn meta_mut(&mut self) -> &mut Option { + fn meta_mut(&mut self) -> &mut Option { &mut self.meta } } @@ -470,7 +738,7 @@ mod tests { #[test] fn trace_context_round_trip() { - let mut meta = Meta::new(); + let mut meta = MetaObject::new(); meta.set_traceparent(TRACEPARENT); meta.set_tracestate("vendor1=value1,vendor2=value2"); meta.set_baggage("userId=alice,region=us-east-1"); @@ -481,7 +749,7 @@ mod tests { #[test] fn absent_field_is_none() { - let meta = Meta::new(); + let meta = MetaObject::new(); assert_eq!(meta.get_traceparent(), None); assert_eq!(meta.get_tracestate(), None); assert_eq!(meta.get_baggage(), None); @@ -489,9 +757,9 @@ mod tests { #[test] fn non_string_value_is_none() { - let mut meta = Meta::new(); + let mut meta = MetaObject::new(); meta.0 - .insert(Meta::TRACEPARENT_FIELD.to_string(), Value::from(42)); + .insert(MetaObject::TRACEPARENT_FIELD.to_string(), Value::from(42)); assert_eq!(meta.get_traceparent(), None); } @@ -502,4 +770,90 @@ mod tests { params.set_traceparent(TRACEPARENT); assert_eq!(params.traceparent(), Some(TRACEPARENT)); } + + #[test] + fn request_meta_derefs_to_general_helpers() { + let mut meta = RequestMetaObject::new(); + meta.set_traceparent(TRACEPARENT); + meta.set_progress_token(ProgressToken(NumberOrString::Number(7))); + assert_eq!(meta.get_traceparent(), Some(TRACEPARENT)); + assert_eq!( + meta.get_progress_token(), + Some(ProgressToken(NumberOrString::Number(7))) + ); + } + + mod subscription_id { + use super::*; + + #[test] + fn returns_none_when_absent() { + let meta = NotificationMetaObject::new(); + assert_eq!(meta.subscription_id(), None); + } + + #[test] + fn round_trips_number_id() { + let mut meta = NotificationMetaObject::new(); + meta.set_subscription_id(RequestId::Number(42)); + assert_eq!(meta.subscription_id(), Some(RequestId::Number(42))); + } + + #[test] + fn round_trips_string_id() { + let mut meta = NotificationMetaObject::new(); + meta.set_subscription_id(RequestId::String("sub-1".into())); + assert_eq!( + meta.subscription_id(), + Some(RequestId::String("sub-1".into())) + ); + } + } + + mod missing_required_keys { + use super::*; + + #[test] + fn is_empty_for_pre_draft_protocols() { + let meta = RequestMetaObject::new(); + assert!( + meta.missing_required_keys(&ProtocolVersion::V_2025_11_25) + .is_empty() + ); + } + + #[test] + fn lists_all_draft_keys_for_empty_meta() { + let meta = RequestMetaObject::new(); + assert_eq!( + meta.missing_required_keys(&ProtocolVersion::V_2026_07_28), + RequestMetaObject::DRAFT_REQUIRED_KEYS.to_vec() + ); + } + + #[test] + fn treats_malformed_values_as_missing() { + let meta: RequestMetaObject = serde_json::from_value(serde_json::json!({ + "io.modelcontextprotocol/protocolVersion": 123, + "io.modelcontextprotocol/clientCapabilities": null, + })) + .unwrap(); + assert_eq!( + meta.missing_required_keys(&ProtocolVersion::V_2026_07_28), + RequestMetaObject::DRAFT_REQUIRED_KEYS.to_vec() + ); + } + + #[test] + fn is_empty_when_draft_keys_are_present() { + let mut meta = RequestMetaObject::new(); + meta.set_protocol_version(ProtocolVersion::V_2026_07_28); + meta.set_client_info(Implementation::from_build_env()); + meta.set_client_capabilities(ClientCapabilities::default()); + assert!( + meta.missing_required_keys(&ProtocolVersion::V_2026_07_28) + .is_empty() + ); + } + } } diff --git a/crates/rmcp/src/model/mrtr.rs b/crates/rmcp/src/model/mrtr.rs new file mode 100644 index 0000000..8dac9a3 --- /dev/null +++ b/crates/rmcp/src/model/mrtr.rs @@ -0,0 +1,558 @@ +//! Multi Round-Trip Request (MRTR) types for SEP-2322. +//! +//! Provides [`InputRequiredResult`], [`InputRequests`], and [`InputResponses`] +//! for the stateless multi round-trip request pattern defined in the MCP spec. +//! [`ResultType`] lives in the parent [`super`] module alongside other base result types. +//! +//! # Overview +//! +//! A server may respond to `tools/call`, `prompts/get`, or `resources/read` with an +//! [`InputRequiredResult`] instead of the normal result. The client fulfills the +//! [`InputRequests`], then retries the original request with [`InputResponses`] and +//! the echoed `requestState`. +//! +//! # Using MRTR +//! +//! **Server:** return an [`InputRequiredResult`] from a tool/prompt/resource +//! handler via the matching outcome enum ([`CallToolResponse`], +//! [`GetPromptResponse`], [`ReadResourceResponse`]). The SDK only lets an +//! `InputRequiredResult` reach a peer that negotiated protocol version +//! `2026-07-28` or newer; older peers get a protocol error instead. +//! +//! **Client:** the high-level `RunningService` helpers — `call_tool`, +//! `get_prompt`, and `read_resource` — automatically fulfil each +//! [`InputRequest`] through the local `ClientHandler` and retry, up to +//! [`DEFAULT_MRTR_MAX_ROUNDS`]. Use the `*_once` variants (e.g. +//! `call_tool_once`) to receive an [`InputRequiredResult`] directly and drive +//! the rounds yourself. +//! +//! # `requestState` is untrusted +//! +//! The client echoes `requestState` back verbatim, so a stateless server that +//! stores meaningful data in it MUST verify integrity before trusting the echoed +//! value. Enable the `request-state` feature and use `RequestStateCodec` to seal +//! and open it, or keep the state server-side and use `requestState` only as an +//! opaque handle. +//! +//! A complete runnable walkthrough lives in the `servers_mrtr` example. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use super::{ + CallToolResult, CreateMessageRequest, CreateTaskResult, ElicitRequest, GetPromptResult, + ListRootsRequest, MetaObject, ReadResourceResult, ResultType, ServerResult, +}; + +/// Default maximum number of MRTR rounds a high-level client call will drive. +/// +/// This matches the default used by other Tier 1 SDKs and prevents a +/// misbehaving peer from keeping a request alive indefinitely. +pub const DEFAULT_MRTR_MAX_ROUNDS: usize = 10; + +/// A server-initiated request that can appear inside [`InputRequests`]. +/// +/// Per the MCP spec, only `CreateMessageRequest` (sampling), +/// `ElicitRequest` (elicitation), and `ListRootsRequest` (roots) +/// are allowed. This is modeled as an untagged enum rather than a +/// `ServerRequest` alias to prevent `PingRequest` or `CustomRequest` from +/// being included. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[non_exhaustive] +pub enum InputRequest { + /// A `sampling/createMessage` request. + CreateMessage(CreateMessageRequest), + /// An `elicitation/create` request. + Elicitation(ElicitRequest), + /// A `roots/list` request. + ListRoots(ListRootsRequest), +} + +// Wire-level equality: two `InputRequest`s are equal if they serialize to the +// same JSON. The wrapped request envelopes do not implement `PartialEq` +// structurally (they carry `Extensions`). +impl PartialEq for InputRequest { + fn eq(&self, other: &Self) -> bool { + match (serde_json::to_value(self), serde_json::to_value(other)) { + (Ok(a), Ok(b)) => a == b, + _ => false, + } + } +} + +/// A map of server-initiated requests that the client must fulfill. +/// +/// Keys are server-assigned string identifiers; values are request objects +/// (`ElicitRequest`, `CreateMessageRequest`, or `ListRootsRequest`). +pub type InputRequests = BTreeMap; + +/// A map of client responses to server-initiated requests. +/// +/// Keys correspond to the keys in the [`InputRequests`] map; values are the +/// client's result for each request (`ElicitResult`, `CreateMessageResult`, +/// or `ListRootsResult`), represented as opaque JSON because the +/// heterogeneous `ClientResult` union does not derive the traits required +/// for use as a `BTreeMap` value. +pub type InputResponses = BTreeMap; + +/// Result of a `tools/call` request, including the MRTR intermediate result. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub enum CallToolResponse { + /// The server completed the tool call. + Complete(CallToolResult), + /// The server requires client-side input before the tool call can complete. + InputRequired(InputRequiredResult), + /// The server materialized a task for this call (SEP-2663 Tasks extension, + /// `resultType: "task"`). The client polls `tasks/get` for the result. + Task(CreateTaskResult), +} + +impl From for CallToolResponse { + fn from(result: CallToolResult) -> Self { + Self::Complete(result) + } +} + +impl From for CallToolResponse { + fn from(result: InputRequiredResult) -> Self { + Self::InputRequired(result) + } +} + +impl From for ServerResult { + fn from(response: CallToolResponse) -> Self { + match response { + CallToolResponse::Complete(result) => ServerResult::CallToolResult(result), + CallToolResponse::InputRequired(result) => ServerResult::InputRequiredResult(result), + CallToolResponse::Task(result) => ServerResult::CreateTaskResult(result), + } + } +} + +impl From for CallToolResponse { + fn from(result: CreateTaskResult) -> Self { + Self::Task(result) + } +} + +/// Result of a `prompts/get` request, including the MRTR intermediate result. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub enum GetPromptResponse { + /// The server completed the prompt request. + Complete(GetPromptResult), + /// The server requires client-side input before the prompt can be returned. + InputRequired(InputRequiredResult), +} + +impl From for GetPromptResponse { + fn from(result: GetPromptResult) -> Self { + Self::Complete(result) + } +} + +impl From for GetPromptResponse { + fn from(result: InputRequiredResult) -> Self { + Self::InputRequired(result) + } +} + +impl From for ServerResult { + fn from(response: GetPromptResponse) -> Self { + match response { + GetPromptResponse::Complete(result) => ServerResult::GetPromptResult(result), + GetPromptResponse::InputRequired(result) => ServerResult::InputRequiredResult(result), + } + } +} + +/// Result of a `resources/read` request, including the MRTR intermediate result. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub enum ReadResourceResponse { + /// The server completed the resource read. + Complete(ReadResourceResult), + /// The server requires client-side input before the resource can be returned. + InputRequired(InputRequiredResult), +} + +impl From for ReadResourceResponse { + fn from(result: ReadResourceResult) -> Self { + Self::Complete(result) + } +} + +impl From for ReadResourceResponse { + fn from(result: InputRequiredResult) -> Self { + Self::InputRequired(result) + } +} + +impl From for ServerResult { + fn from(response: ReadResourceResponse) -> Self { + match response { + ReadResourceResponse::Complete(result) => ServerResult::ReadResourceResult(result), + ReadResourceResponse::InputRequired(result) => { + ServerResult::InputRequiredResult(result) + } + } + } +} + +/// A result indicating that additional input is needed before the request +/// can be completed. +/// +/// At least one of [`input_requests`](Self::input_requests) or +/// [`request_state`](Self::request_state) MUST be present. +/// +/// Servers MAY send this in response to `tools/call`, `prompts/get`, or +/// `resources/read`. Servers MUST NOT send this for any other request. +/// +/// # Examples +/// +/// ``` +/// use rmcp::model::InputRequiredResult; +/// +/// let result = InputRequiredResult::from_request_state("opaque-server-state"); +/// assert!(result.input_requests.is_none()); +/// assert_eq!(result.request_state.as_deref(), Some("opaque-server-state")); +/// ``` +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[non_exhaustive] +pub struct InputRequiredResult { + /// Always `"input_required"` for this result type. + pub result_type: ResultType, + + /// Server-initiated requests that the client must fulfill before retrying. + #[serde(skip_serializing_if = "Option::is_none")] + pub input_requests: Option, + + /// Opaque request state to be echoed back by the client on retry. + /// Clients MUST NOT inspect, parse, modify, or make any assumptions + /// about the contents. + #[serde(skip_serializing_if = "Option::is_none")] + pub request_state: Option, + + /// Optional protocol-level metadata. + #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] + pub meta: Option, +} + +/// Custom deserializer that requires `resultType: "input_required"` to prevent +/// greedy matching in the untagged `ServerResult` enum (which would otherwise +/// swallow empty objects or unknown shapes). +impl<'de> Deserialize<'de> for InputRequiredResult { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + #[derive(Deserialize)] + #[serde(rename_all = "camelCase")] + struct Helper { + result_type: Option, + input_requests: Option, + request_state: Option, + #[serde(rename = "_meta")] + meta: Option, + } + + let helper = Helper::deserialize(deserializer)?; + + match &helper.result_type { + Some(rt) if rt.is_input_required() => {} + _ => { + return Err(serde::de::Error::custom( + "InputRequiredResult requires resultType to be \"input_required\"", + )); + } + } + + if helper.input_requests.is_none() && helper.request_state.is_none() { + return Err(serde::de::Error::custom( + "InputRequiredResult requires at least one of inputRequests or requestState", + )); + } + + Ok(InputRequiredResult { + result_type: ResultType::INPUT_REQUIRED, + input_requests: helper.input_requests, + request_state: helper.request_state, + meta: helper.meta, + }) + } +} + +impl InputRequiredResult { + /// Creates a new `InputRequiredResult` with both input requests and request state. + pub fn new(input_requests: Option, request_state: Option) -> Self { + Self { + result_type: ResultType::INPUT_REQUIRED, + input_requests, + request_state, + meta: None, + } + } + + /// Creates from input requests only. + pub fn from_input_requests(input_requests: InputRequests) -> Self { + Self::new(Some(input_requests), None) + } + + /// Creates from request state only (e.g. for load shedding). + pub fn from_request_state(request_state: impl Into) -> Self { + Self::new(None, Some(request_state.into())) + } + + /// Sets optional metadata. + pub fn with_meta(mut self, meta: MetaObject) -> Self { + self.meta = Some(meta); + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + + mod result_type { + use super::*; + + #[test] + fn default_is_complete() { + assert_eq!(ResultType::default(), ResultType::COMPLETE); + } + + #[test] + fn serializes_complete() { + assert_eq!( + serde_json::to_value(&ResultType::COMPLETE).unwrap(), + serde_json::json!("complete") + ); + } + + #[test] + fn serializes_input_required() { + assert_eq!( + serde_json::to_value(&ResultType::INPUT_REQUIRED).unwrap(), + serde_json::json!("input_required") + ); + } + + #[test] + fn deserializes_known_values() { + let complete: ResultType = + serde_json::from_value(serde_json::json!("complete")).unwrap(); + assert_eq!(complete, ResultType::COMPLETE); + + let input_required: ResultType = + serde_json::from_value(serde_json::json!("input_required")).unwrap(); + assert_eq!(input_required, ResultType::INPUT_REQUIRED); + } + + #[test] + fn preserves_unknown_extension_values() { + let custom: ResultType = + serde_json::from_value(serde_json::json!("streaming")).unwrap(); + assert_eq!(custom.as_str(), "streaming"); + assert!(!custom.is_complete()); + assert!(!custom.is_input_required()); + + let reserialized = serde_json::to_value(&custom).unwrap(); + assert_eq!(reserialized, serde_json::json!("streaming")); + } + } + + mod input_required_result { + use super::*; + + #[test] + fn deserializes_with_requests_and_state() { + let json = serde_json::json!({ + "resultType": "input_required", + "inputRequests": { + "github_login": { + "method": "elicitation/create", + "params": { + "message": "Please provide your GitHub username", + "requestedSchema": { + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"] + } + } + }, + "capital_of_france": { + "method": "sampling/createMessage", + "params": { + "messages": [{ + "role": "user", + "content": { "type": "text", "text": "What is the capital of France?" } + }], + "maxTokens": 100 + } + } + }, + "requestState": "eyJsb2NhdGlvbiI6Ik5ldyBZb3JrIn0" + }); + + let result: InputRequiredResult = serde_json::from_value(json).unwrap(); + + let requests = result + .input_requests + .as_ref() + .expect("should have input_requests"); + assert_eq!(requests.len(), 2); + assert!(requests.contains_key("github_login")); + assert!(requests.contains_key("capital_of_france")); + assert_eq!( + result.request_state.as_deref(), + Some("eyJsb2NhdGlvbiI6Ik5ldyBZb3JrIn0") + ); + } + + #[test] + fn roundtrip_preserves_all_fields() { + let json = serde_json::json!({ + "resultType": "input_required", + "inputRequests": { + "key": { + "method": "elicitation/create", + "params": { + "message": "test", + "requestedSchema": { "type": "object", "properties": {} } + } + } + }, + "requestState": "abc123" + }); + + let result: InputRequiredResult = serde_json::from_value(json).unwrap(); + let reserialized = serde_json::to_value(&result).unwrap(); + + assert_eq!(reserialized["resultType"], "input_required"); + assert!(reserialized["inputRequests"].is_object()); + assert_eq!(reserialized["requestState"], "abc123"); + } + + #[test] + fn deserializes_with_request_state_only() { + let json = serde_json::json!({ + "resultType": "input_required", + "requestState": "eyJwcm9ncmVzcyI6IjUwJSJ9" + }); + + let result: InputRequiredResult = serde_json::from_value(json).unwrap(); + + assert!(result.input_requests.is_none()); + assert_eq!( + result.request_state.as_deref(), + Some("eyJwcm9ncmVzcyI6IjUwJSJ9") + ); + } + + #[test] + fn rejects_missing_input_requests_and_request_state() { + let json = serde_json::json!({ + "resultType": "input_required", + "_meta": {} + }); + let err = serde_json::from_value::(json).unwrap_err(); + assert!( + err.to_string().contains("inputRequests or requestState"), + "error should mention the required continuation fields, got: {err}" + ); + } + + #[test] + fn rejects_missing_result_type() { + let json = serde_json::json!({ + "requestState": "some-state" + }); + let err = serde_json::from_value::(json).unwrap_err(); + assert!( + err.to_string().contains("input_required"), + "error should mention the required resultType, got: {err}" + ); + } + + #[test] + fn rejects_wrong_result_type() { + let json = serde_json::json!({ + "resultType": "complete", + "requestState": "some-state" + }); + let err = serde_json::from_value::(json).unwrap_err(); + assert!( + err.to_string().contains("input_required"), + "error should mention the required resultType, got: {err}" + ); + } + } + + mod input_responses { + use super::*; + + #[test] + fn deserializes_heterogeneous_results() { + let json = serde_json::json!({ + "github_login": { + "action": "accept", + "content": { "name": "octocat" } + }, + "capital_of_france": { + "role": "assistant", + "content": { "type": "text", "text": "Paris." }, + "model": "claude-3-sonnet-20240307", + "stopReason": "endTurn" + } + }); + + let responses: InputResponses = serde_json::from_value(json).unwrap(); + + assert_eq!(responses.len(), 2); + assert!(responses.contains_key("github_login")); + assert!(responses.contains_key("capital_of_france")); + } + } + + mod constructors { + use super::*; + + #[test] + fn from_request_state_sets_state_only() { + let result = InputRequiredResult::from_request_state("opaque"); + + assert_eq!(result.result_type, ResultType::INPUT_REQUIRED); + assert!(result.input_requests.is_none()); + assert_eq!(result.request_state.as_deref(), Some("opaque")); + } + + #[test] + fn from_input_requests_sets_requests_only() { + let mut requests = InputRequests::new(); + requests.insert( + "key".to_string(), + serde_json::from_value(serde_json::json!({ + "method": "elicitation/create", + "params": { + "message": "test", + "requestedSchema": { "type": "object", "properties": {} } + } + })) + .unwrap(), + ); + + let result = InputRequiredResult::from_input_requests(requests); + + assert!(result.input_requests.is_some()); + assert!(result.request_state.is_none()); + } + } +} diff --git a/crates/rmcp/src/model/prompt.rs b/crates/rmcp/src/model/prompt.rs index e438260..46f5168 100644 --- a/crates/rmcp/src/model/prompt.rs +++ b/crates/rmcp/src/model/prompt.rs @@ -1,7 +1,7 @@ use serde::{Deserialize, Serialize}; use super::{ - Annotations, ContentBlock, Icon, Meta, Role, + Annotations, ContentBlock, Icon, MetaObject, Role, content::{AudioContent, EmbeddedResource, ImageContent, TextContent}, resource::ResourceContents, }; @@ -22,7 +22,7 @@ pub struct Prompt { #[serde(skip_serializing_if = "Option::is_none")] pub icons: Option>, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, } impl Prompt { @@ -70,7 +70,7 @@ impl Prompt { self } - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } @@ -144,7 +144,7 @@ impl PromptMessage { role: Role, data: &[u8], mime_type: &str, - meta: Option, + meta: Option, annotations: Option, ) -> Self { use base64::{Engine, prelude::BASE64_STANDARD}; @@ -166,7 +166,7 @@ impl PromptMessage { role: Role, data: &[u8], mime_type: &str, - meta: Option, + meta: Option, annotations: Option, ) -> Self { use base64::{Engine, prelude::BASE64_STANDARD}; @@ -188,8 +188,8 @@ impl PromptMessage { uri: String, mime_type: Option, text: Option, - resource_meta: Option, - resource_content_meta: Option, + resource_meta: Option, + resource_content_meta: Option, annotations: Option, ) -> Self { let resource_contents = match text { @@ -216,7 +216,11 @@ impl PromptMessage { } } - pub fn new_text_with_meta>(role: Role, text: S, meta: Option) -> Self { + pub fn new_text_with_meta>( + role: Role, + text: S, + meta: Option, + ) -> Self { Self { role, content: ContentBlock::Text(TextContent { diff --git a/crates/rmcp/src/model/request_state.rs b/crates/rmcp/src/model/request_state.rs new file mode 100644 index 0000000..9f45572 --- /dev/null +++ b/crates/rmcp/src/model/request_state.rs @@ -0,0 +1,1494 @@ +//! Integrity protection for SEP-2322 `requestState`. +//! +//! In the multi round-trip request (MRTR) flow, a server places an opaque +//! `requestState` string in an [`InputRequiredResult`](super::InputRequiredResult) +//! and the client echoes it back verbatim on retry. From the server's point of +//! view the echoed value is **untrusted, attacker-controlled input**: a client +//! can send back anything it likes. Per SEP-2322, a server that lets +//! `requestState` influence authorization, resource access, or business logic +//! MUST protect its integrity and reject values that fail verification. +//! +//! [`RequestStateCodec`] provides an opt-in way to do this. It seals a payload +//! into an opaque string with an HMAC-SHA256 tag and opens it again, rejecting +//! any value that was forged or tampered with. +//! +//! To follow the spec's replay-prevention guidance without hand-rolling the +//! checks, the codec supports two bindings via [`SealOptions`]: +//! +//! * **Associated data** — arbitrary context (e.g. the authenticated principal +//! plus a digest of the originating request) that is mixed into the tag but +//! not stored in the token. [`open_with`](RequestStateCodec::open_with) only +//! succeeds when the caller supplies the same context, so a value cannot be +//! replayed by a different principal or against a different request. This is +//! *fail-closed*: forgetting to pass the context makes verification fail. +//! * **TTL** — a relative expiry stamped into the token; opening a value past +//! its expiry fails with [`RequestStateError::Expired`]. +//! +//! Single-use/nonce enforcement (for one-time redemptions) still has to be done +//! server-side, as the spec notes. +//! +//! This helper is only about *integrity*, not *confidentiality*: the sealed +//! payload is signed, not encrypted, so it is base64url-readable by anyone. Do +//! not put secrets in it. +//! +//! Using the codec is entirely optional. A server that keeps its state +//! server-side, or that does not trust `requestState` for anything security +//! sensitive, can keep building the string by hand via +//! [`InputRequiredResult::from_request_state`](super::InputRequiredResult::from_request_state). +//! +//! # Examples +//! +//! ``` +//! use rmcp::model::{RequestStateCodec, SealOptions}; +//! # fn main() -> Result<(), rmcp::model::RequestStateError> { +//! +//! // Derive the key from a per-process secret; keep it out of client reach. +//! let codec = RequestStateCodec::try_new(b"a-32-byte-or-longer-secret-key!!!")?; +//! +//! // Bind the state to the caller and the originating request. +//! let context = b"user:alice|tools/call:weather"; +//! let sealed = codec.seal_with( +//! b"step=2", +//! &SealOptions::new().associated_data(context), +//! ); +//! +//! // On retry the client echoes `sealed` back untouched; the server re-derives +//! // the same context and opens it. +//! let opened = codec.open_with(&sealed, context)?; +//! assert_eq!(opened, b"step=2"); +//! +//! // A different principal (different context) is rejected. +//! assert!(codec.open_with(&sealed, b"user:bob|tools/call:weather").is_err()); +//! # Ok(()) +//! # } +//! ``` + +use std::{ + collections::{HashMap, hash_map::Entry}, + time::Duration, +}; + +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; +use hmac::{Hmac, KeyInit, Mac}; +use serde::{Serialize, de::DeserializeOwned}; +use sha2::Sha256; +use thiserror::Error; +use zeroize::Zeroizing; + +type HmacSha256 = Hmac; + +const VERSION_V1: &str = "rs1"; + +const VERSION_V2: &str = "rs2"; + +/// Domain-separation label mixed into the HMAC so a `requestState` tag can never +/// be confused with an HMAC computed for some other purpose using the same key. +const DOMAIN_V1: &[u8] = b"rmcp/mrtr/request-state/v1"; + +/// Domain-separation label for keyed request-state tags. +const DOMAIN_V2: &[u8] = b"rmcp/mrtr/request-state/v2"; + +/// Length of the big-endian expiry prefix (unix milliseconds) stored at the +/// front of every sealed body. `0` means "no expiry". +const EXPIRY_LEN: usize = 8; + +const MAX_KID_LEN: usize = 255; + +/// Maximum unpadded base64url length for [`MAX_KID_LEN`] bytes. +const MAX_ENCODED_KID_LEN: usize = 340; + +/// Errors returned when constructing or configuring a +/// [`RequestStateCodec`], or when processing a sealed value. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum RequestStateError { + /// The signing key is shorter than [`RequestStateCodec::MIN_KEY_LENGTH`]. + #[error( + "request state signing key is too short: expected at least {minimum} bytes, got {actual}" + )] + KeyTooShort { minimum: usize, actual: usize }, + + /// The value is not a well-formed sealed request state (wrong prefix or + /// missing sections). + #[error("request state is malformed or uses an unsupported format")] + MalformedFormat, + + /// A section of the value was not valid base64url. + #[error("request state is not valid base64url")] + InvalidEncoding, + + /// The HMAC tag did not match; the value was forged, tampered with, or + /// opened with the wrong associated data. + #[error("request state failed integrity verification")] + IntegrityCheckFailed, + + /// The value carried a TTL that has already elapsed. + #[error("request state has expired")] + Expired, + + /// The token names (or, for rs1, requires) a key this codec does not hold. + #[error("request state was sealed with an unknown key")] + UnknownKeyId, + + /// The decoded rs2 key id was empty, too long, or not valid UTF-8. + #[error("request state contains an invalid key identifier")] + InvalidKeyId, + + /// An invalid keyring configuration; the message is diagnostic only. + #[error("invalid keyring configuration: {0}")] + InvalidKeyring(&'static str), + + /// The sealed payload could not be serialized to JSON. + #[error("failed to serialize request state payload: {0}")] + Serialization(#[source] serde_json::Error), + + /// The opened payload could not be deserialized from JSON. + #[error("failed to deserialize request state payload: {0}")] + Deserialization(#[source] serde_json::Error), +} + +/// Options controlling how a value is sealed by [`RequestStateCodec`]. +/// +/// Defaults to no associated data and no expiry, which is equivalent to the +/// bare [`seal`](RequestStateCodec::seal) / [`open`](RequestStateCodec::open) +/// methods. +#[derive(Clone, Copy, Debug, Default)] +pub struct SealOptions<'a> { + associated_data: &'a [u8], + ttl: Option, +} + +impl<'a> SealOptions<'a> { + /// Creates empty options (no associated data, no expiry). + pub fn new() -> Self { + Self::default() + } + + /// Binds the sealed value to `associated_data`. The same bytes must be + /// supplied to [`open_with`](RequestStateCodec::open_with); the data is + /// authenticated but not stored in the token. + /// + /// Use this to bind the state to the authenticated principal and/or the + /// originating request (e.g. method name plus a digest of its parameters). + pub fn associated_data(mut self, associated_data: &'a [u8]) -> Self { + self.associated_data = associated_data; + self + } + + /// Sets a relative time-to-live after which opening the value fails with + /// [`RequestStateError::Expired`]. + pub fn ttl(mut self, ttl: Duration) -> Self { + self.ttl = Some(ttl); + self + } +} + +/// A keyed codec for integrity-protected SEP-2322 `requestState` values. +/// +/// [`new`](Self::new) preserves `rs1`; [`new_with_keyring`](Self::new_with_keyring) +/// emits `rs2`. +/// Use [`with_rs1_signing`](Self::with_rs1_signing) and +/// [`with_rs1_fallback`](Self::with_rs1_fallback) for rolling migrations. +/// +/// Use [`try_new`](Self::try_new) and +/// [`new_with_keyring`](Self::new_with_keyring) to require at least +/// [`MIN_KEY_LENGTH`](Self::MIN_KEY_LENGTH) bytes of high-entropy key material. +/// Configure the same keys on every replica. Stored key material is zeroized +/// when the codec is dropped. +/// +/// Values are authenticated, not encrypted; key ids and payloads remain +/// readable. Retain old keys until every value they signed has expired. +/// +/// # Key rotation +/// +/// Rotate in stages so every replica can open states emitted by its peers: +/// +/// ``` +/// # use rmcp::model::{RequestStateCodec, RequestStateError}; +/// # fn main() -> Result<(), RequestStateError> { +/// let old_key = b"old-request-state-key-at-least-32b".as_slice(); +/// let new_key = b"new-request-state-key-at-least-32b".as_slice(); +/// let keys = [("old", old_key), ("new", new_key)]; +/// +/// // 1. Deploy both keys, but continue emitting rs1 with the old key. +/// let transitional = RequestStateCodec::new_with_keyring("new", keys)? +/// .with_rs1_signing("old")?; +/// +/// // 2. Emit rs2 with the new key, while accepting in-flight rs1 values. +/// let promoted = RequestStateCodec::new_with_keyring("new", keys)? +/// .with_rs1_fallback("old")?; +/// +/// let old_state = transitional.seal(b"old state"); +/// assert_eq!(promoted.open(&old_state)?, b"old state"); +/// let new_state = promoted.seal(b"new state"); +/// assert_eq!(transitional.open(&new_state)?, b"new state"); +/// +/// // 3. After old values expire, remove the old key. +/// let retired = RequestStateCodec::new_with_keyring("new", [("new", new_key)])?; +/// assert_eq!(retired.open(&new_state)?, b"new state"); +/// # Ok(()) +/// # } +/// ``` +#[derive(Clone)] +pub struct RequestStateCodec { + keys: Keys, +} + +#[derive(Clone)] +enum Keys { + Single(Zeroizing>), + Ring { + keys: HashMap>>, + seal_mode: SealMode, + rs1_fallbacks: Vec, + }, +} + +#[derive(Clone, Debug)] +enum SealMode { + Rs1 { key_id: String }, + Rs2 { key_id: String }, +} + +impl std::fmt::Debug for RequestStateCodec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Never leak signing or verification keys through Debug output. + match &self.keys { + Keys::Single(_) => f + .debug_struct("RequestStateCodec") + .field("mode", &"single") + .field("key", &"") + .finish(), + Keys::Ring { + keys, + seal_mode, + rs1_fallbacks, + } => f + .debug_struct("RequestStateCodec") + .field("mode", &"ring") + .field("key_count", &keys.len()) + .field("keys", &"") + .field("seal_mode", seal_mode) + .field("rs1_fallbacks", rs1_fallbacks) + .finish(), + } + } +} + +impl RequestStateCodec { + /// Minimum accepted signing-key length in bytes. + pub const MIN_KEY_LENGTH: usize = 32; + + /// Creates a legacy single-key codec that seals and opens `rs1` values + /// without validating the signing-key length. + #[deprecated( + since = "3.1.4", + note = "use RequestStateCodec::try_new to enforce the minimum signing-key length" + )] + pub fn new(key: impl Into>) -> Self { + Self::new_unchecked(key) + } + + /// Creates a codec without validating the signing-key length. + /// + /// Prefer [`try_new`](Self::try_new) unless the key length has already been + /// validated. + pub fn new_unchecked(key: impl Into>) -> Self { + Self { + keys: Keys::Single(Zeroizing::new(key.into())), + } + } + + /// Creates a legacy single-key codec after validating its signing-key length. + /// + /// # Errors + /// + /// Returns [`RequestStateError::KeyTooShort`] when `key` contains fewer + /// than [`MIN_KEY_LENGTH`](Self::MIN_KEY_LENGTH) bytes. + pub fn try_new(key: impl Into>) -> Result { + let key = Zeroizing::new(key.into()); + Self::validate_key_length(&key)?; + Ok(Self { + keys: Keys::Single(key), + }) + } + + /// Creates a keyring codec that seals `rs2` with `active_kid`. + /// + /// All configured keys can open matching `rs2` values. Legacy `rs1` values + /// require [`with_rs1_fallback`](Self::with_rs1_fallback) or + /// [`with_rs1_signing`](Self::with_rs1_signing). + /// Key ids are case-sensitive UTF-8 strings of 1 to 255 bytes. + /// + /// # Errors + /// + /// Returns [`RequestStateError::KeyTooShort`] if a configured key contains + /// fewer than [`MIN_KEY_LENGTH`](Self::MIN_KEY_LENGTH) bytes. Returns + /// [`RequestStateError::InvalidKeyring`] if the keyring is empty, contains + /// duplicate or invalid key ids, or does not contain `active_kid`. + pub fn new_with_keyring( + active_kid: impl Into, + keys: impl IntoIterator, + ) -> Result + where + K: Into, + V: Into>, + { + let mut ring = HashMap::new(); + for (kid, key) in keys { + let kid = kid.into(); + let key = Zeroizing::new(key.into()); + Self::validate_config_kid(&kid)?; + Self::validate_key_length(&key)?; + match ring.entry(kid) { + Entry::Vacant(entry) => { + entry.insert(key); + } + Entry::Occupied(_) => { + return Err(RequestStateError::InvalidKeyring("duplicate key id")); + } + } + } + + if ring.is_empty() { + return Err(RequestStateError::InvalidKeyring( + "keyring must contain at least one key", + )); + } + + let active_kid = active_kid.into(); + Self::validate_config_kid(&active_kid)?; + if !ring.contains_key(&active_kid) { + return Err(RequestStateError::InvalidKeyring( + "active key id is not present in the keyring", + )); + } + + Ok(Self { + keys: Keys::Ring { + keys: ring, + seal_mode: SealMode::Rs2 { key_id: active_kid }, + rs1_fallbacks: Vec::new(), + }, + }) + } + + /// Uses `kid` to sign legacy `rs1` values during a rolling migration. + /// The codec accepts its own `rs1` output and configured `rs2` values. + /// + /// # Errors + /// + /// Returns [`RequestStateError::InvalidKeyring`] when called on a single-key + /// codec, or if the keyring does not contain `kid`. + pub fn with_rs1_signing(mut self, kid: impl AsRef) -> Result { + let kid = kid.as_ref(); + match &mut self.keys { + Keys::Single(_) => Err(RequestStateError::InvalidKeyring( + "rs1 transitional signing requires a keyring", + )), + Keys::Ring { + keys, + seal_mode, + rs1_fallbacks, + } => { + if !keys.contains_key(kid) { + return Err(RequestStateError::InvalidKeyring( + "rs1 signing key id is not present in the keyring", + )); + } + + let kid = kid.to_owned(); + *seal_mode = SealMode::Rs1 { + key_id: kid.clone(), + }; + if !rs1_fallbacks.iter().any(|existing| existing == &kid) { + rs1_fallbacks.push(kid); + } + Ok(self) + } + } + } + + /// Accepts legacy `rs1` values signed with `kid`. + /// Adding the same id more than once has no effect; remove fallbacks after + /// old values have drained. + /// + /// # Errors + /// + /// Returns [`RequestStateError::InvalidKeyring`] when called on a single-key + /// codec, or if the keyring does not contain `kid`. + pub fn with_rs1_fallback(mut self, kid: impl AsRef) -> Result { + let kid = kid.as_ref(); + match &mut self.keys { + Keys::Single(_) => Err(RequestStateError::InvalidKeyring( + "rs1 fallbacks require a keyring", + )), + Keys::Ring { + keys, + rs1_fallbacks, + .. + } => { + if !keys.contains_key(kid) { + return Err(RequestStateError::InvalidKeyring( + "rs1 fallback key id is not present in the keyring", + )); + } + if !rs1_fallbacks.iter().any(|existing| existing == kid) { + rs1_fallbacks.push(kid.to_owned()); + } + Ok(self) + } + } + } + + /// Seals raw bytes into an opaque, integrity-protected string suitable for + /// use as `requestState`. + pub fn seal(&self, payload: &[u8]) -> String { + self.seal_with(payload, &SealOptions::default()) + } + + /// Seals raw bytes with [`SealOptions`] (associated data and/or TTL). + pub fn seal_with(&self, payload: &[u8], options: &SealOptions<'_>) -> String { + self.seal_at(payload, options, Self::now_ms()) + } + + /// Seals a serializable value by encoding it as JSON before sealing. + /// + /// # Errors + /// + /// Returns [`RequestStateError::Serialization`] if `value` cannot be encoded + /// as JSON. + pub fn seal_json(&self, value: &T) -> Result { + self.seal_json_with(value, &SealOptions::default()) + } + + /// Seals a serializable value with [`SealOptions`]. + /// + /// # Errors + /// + /// Returns [`RequestStateError::Serialization`] if `value` cannot be encoded + /// as JSON. + pub fn seal_json_with( + &self, + value: &T, + options: &SealOptions<'_>, + ) -> Result { + let payload = serde_json::to_vec(value).map_err(RequestStateError::Serialization)?; + Ok(self.seal_with(&payload, options)) + } + + /// Opens a sealed value that was sealed without associated data, verifying + /// its integrity and expiry and returning the original bytes. + /// + /// # Errors + /// + /// See [`open_with`](Self::open_with). + pub fn open(&self, sealed: &str) -> Result, RequestStateError> { + self.open_with(sealed, &[]) + } + + /// Opens a sealed value, verifying its integrity against `associated_data` + /// and checking its expiry. + /// + /// `associated_data` must match the bytes passed to + /// [`SealOptions::associated_data`] when the value was sealed (use `&[]` for + /// values sealed without it). + /// + /// # Errors + /// + /// Returns a [`RequestStateError`] if the value is malformed, cannot be + /// verified, names an unavailable key, or has expired. + /// + /// Applications MUST map all token-opening failures to a single + /// client-facing error and reserve the detailed variants for internal + /// diagnostics. + pub fn open_with( + &self, + sealed: &str, + associated_data: &[u8], + ) -> Result, RequestStateError> { + self.open_at(sealed, associated_data, Self::now_ms()) + } + + /// Opens a sealed value (no associated data) and deserializes its JSON + /// payload. + /// + /// # Errors + /// + /// See [`open_json_with`](Self::open_json_with). + pub fn open_json(&self, sealed: &str) -> Result { + self.open_json_with(sealed, &[]) + } + + /// Opens a sealed value against `associated_data` and deserializes its JSON + /// payload. + /// + /// # Errors + /// + /// Returns the same integrity, expiry, and format errors as + /// [`open_with`](Self::open_with), plus [`RequestStateError::Deserialization`] + /// if the payload is not valid JSON for `T`. + pub fn open_json_with( + &self, + sealed: &str, + associated_data: &[u8], + ) -> Result { + let payload = self.open_with(sealed, associated_data)?; + serde_json::from_slice(&payload).map_err(RequestStateError::Deserialization) + } + + fn seal_at(&self, payload: &[u8], options: &SealOptions<'_>, now_ms: i64) -> String { + let expiry = match options.ttl { + Some(ttl) => now_ms.saturating_add(ttl.as_millis().min(i64::MAX as u128) as i64), + None => 0, + }; + + // body = big-endian expiry (0 = none) followed by the caller payload. + let mut body = Vec::with_capacity(EXPIRY_LEN + payload.len()); + body.extend_from_slice(&expiry.to_be_bytes()); + body.extend_from_slice(payload); + + match &self.keys { + Keys::Single(key) => Self::seal_rs1(key, options.associated_data, &body), + Keys::Ring { + keys, seal_mode, .. + } => match seal_mode { + SealMode::Rs1 { key_id } => Self::seal_rs1( + keys.get(key_id).expect("validated rs1 signing key"), + options.associated_data, + &body, + ), + SealMode::Rs2 { key_id } => Self::seal_rs2( + key_id, + keys.get(key_id).expect("validated rs2 signing key"), + options.associated_data, + &body, + ), + }, + } + } + + fn open_at( + &self, + sealed: &str, + associated_data: &[u8], + now_ms: i64, + ) -> Result, RequestStateError> { + match sealed.split('.').next() { + Some(VERSION_V1) => self.open_rs1_at(sealed, associated_data, now_ms), + Some(VERSION_V2) => self.open_rs2_at(sealed, associated_data, now_ms), + _ => Err(RequestStateError::MalformedFormat), + } + } + + fn seal_rs1(key: &[u8], associated_data: &[u8], body: &[u8]) -> String { + let tag = Self::mac_v1(key, associated_data, body) + .finalize() + .into_bytes(); + let mut out = String::with_capacity( + VERSION_V1.len() + 2 + Self::b64_len(body.len()) + Self::b64_len(tag.len()), + ); + out.push_str(VERSION_V1); + out.push('.'); + URL_SAFE_NO_PAD.encode_string(body, &mut out); + out.push('.'); + URL_SAFE_NO_PAD.encode_string(tag.as_slice(), &mut out); + out + } + + fn seal_rs2(kid: &str, key: &[u8], associated_data: &[u8], body: &[u8]) -> String { + let tag = Self::mac_v2(key, kid.as_bytes(), associated_data, body) + .finalize() + .into_bytes(); + let mut out = String::with_capacity( + VERSION_V2.len() + + 3 + + Self::b64_len(kid.len()) + + Self::b64_len(body.len()) + + Self::b64_len(tag.len()), + ); + out.push_str(VERSION_V2); + out.push('.'); + URL_SAFE_NO_PAD.encode_string(kid.as_bytes(), &mut out); + out.push('.'); + URL_SAFE_NO_PAD.encode_string(body, &mut out); + out.push('.'); + URL_SAFE_NO_PAD.encode_string(tag.as_slice(), &mut out); + out + } + + fn open_rs1_at( + &self, + sealed: &str, + associated_data: &[u8], + now_ms: i64, + ) -> Result, RequestStateError> { + let mut parts = sealed.split('.'); + let version = parts.next().ok_or(RequestStateError::MalformedFormat)?; + let body_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?; + let tag_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?; + if parts.next().is_some() || version != VERSION_V1 { + return Err(RequestStateError::MalformedFormat); + } + + if matches!( + &self.keys, + Keys::Ring { rs1_fallbacks, .. } if rs1_fallbacks.is_empty() + ) { + return Err(RequestStateError::UnknownKeyId); + } + + let body = URL_SAFE_NO_PAD + .decode(body_b64) + .map_err(|_| RequestStateError::InvalidEncoding)?; + let tag = URL_SAFE_NO_PAD + .decode(tag_b64) + .map_err(|_| RequestStateError::InvalidEncoding)?; + + // Authentication must not short-circuit based on tag bytes. + match &self.keys { + Keys::Single(key) => Self::mac_v1(key, associated_data, &body) + .verify_slice(&tag) + .map_err(|_| RequestStateError::IntegrityCheckFailed)?, + Keys::Ring { + keys, + rs1_fallbacks, + .. + } => { + // Try every fallback so timing does not identify the matching key. + let mut verified = false; + for kid in rs1_fallbacks { + let key = keys.get(kid).expect("validated rs1 fallback key"); + let matches = Self::mac_v1(key, associated_data, &body) + .verify_slice(&tag) + .is_ok(); + verified |= matches; + } + if !verified { + return Err(RequestStateError::IntegrityCheckFailed); + } + } + } + + Self::open_authenticated_body(body, now_ms) + } + + fn open_rs2_at( + &self, + sealed: &str, + associated_data: &[u8], + now_ms: i64, + ) -> Result, RequestStateError> { + let mut parts = sealed.split('.'); + let version = parts.next().ok_or(RequestStateError::MalformedFormat)?; + let kid_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?; + let body_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?; + let tag_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?; + if parts.next().is_some() || version != VERSION_V2 { + return Err(RequestStateError::MalformedFormat); + } + + if kid_b64.len() > MAX_ENCODED_KID_LEN { + return Err(RequestStateError::InvalidKeyId); + } + let kid = URL_SAFE_NO_PAD + .decode(kid_b64) + .map_err(|_| RequestStateError::InvalidEncoding)?; + if kid.is_empty() || kid.len() > MAX_KID_LEN { + return Err(RequestStateError::InvalidKeyId); + } + let kid = String::from_utf8(kid).map_err(|_| RequestStateError::InvalidKeyId)?; + + let key = match &self.keys { + Keys::Single(_) => return Err(RequestStateError::UnknownKeyId), + Keys::Ring { keys, .. } => keys.get(&kid).ok_or(RequestStateError::UnknownKeyId)?, + }; + + let body = URL_SAFE_NO_PAD + .decode(body_b64) + .map_err(|_| RequestStateError::InvalidEncoding)?; + let tag = URL_SAFE_NO_PAD + .decode(tag_b64) + .map_err(|_| RequestStateError::InvalidEncoding)?; + + // Authentication must not short-circuit based on tag bytes. + Self::mac_v2(key, kid.as_bytes(), associated_data, &body) + .verify_slice(&tag) + .map_err(|_| RequestStateError::IntegrityCheckFailed)?; + + Self::open_authenticated_body(body, now_ms) + } + + fn open_authenticated_body(body: Vec, now_ms: i64) -> Result, RequestStateError> { + // The body is now authenticated, so its framing can be trusted. + if body.len() < EXPIRY_LEN { + return Err(RequestStateError::MalformedFormat); + } + let expiry = i64::from_be_bytes(body[..EXPIRY_LEN].try_into().expect("checked length")); + if expiry != 0 && now_ms > expiry { + return Err(RequestStateError::Expired); + } + + Ok(body[EXPIRY_LEN..].to_vec()) + } + + fn mac_v1(key: &[u8], associated_data: &[u8], body: &[u8]) -> HmacSha256 { + let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts keys of any length"); + mac.update(DOMAIN_V1); + mac.update(&(associated_data.len() as u64).to_be_bytes()); + mac.update(associated_data); + mac.update(body); + mac + } + + fn mac_v2(key: &[u8], kid: &[u8], associated_data: &[u8], body: &[u8]) -> HmacSha256 { + let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts keys of any length"); + mac.update(DOMAIN_V2); + mac.update(&(kid.len() as u64).to_be_bytes()); + mac.update(kid); + mac.update(&(associated_data.len() as u64).to_be_bytes()); + mac.update(associated_data); + mac.update(body); + mac + } + + fn validate_key_length(key: &[u8]) -> Result<(), RequestStateError> { + if key.len() < Self::MIN_KEY_LENGTH { + return Err(RequestStateError::KeyTooShort { + minimum: Self::MIN_KEY_LENGTH, + actual: key.len(), + }); + } + Ok(()) + } + + fn validate_config_kid(kid: &str) -> Result<(), RequestStateError> { + if kid.is_empty() { + return Err(RequestStateError::InvalidKeyring( + "key id must not be empty", + )); + } + if kid.len() > MAX_KID_LEN { + return Err(RequestStateError::InvalidKeyring( + "key id exceeds 255 UTF-8 bytes", + )); + } + Ok(()) + } + + // Base64url without padding encodes at most three bytes as four characters. + fn b64_len(len: usize) -> usize { + len.div_ceil(3) * 4 + } + + fn now_ms() -> i64 { + chrono::Utc::now().timestamp_millis() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn try_new_accepts_key_at_minimum_length() { + let result = RequestStateCodec::try_new(vec![0; RequestStateCodec::MIN_KEY_LENGTH]); + + assert!(result.is_ok(), "unexpected result: {result:?}"); + } + + #[test] + fn try_new_rejects_key_below_minimum_length() { + let actual = RequestStateCodec::MIN_KEY_LENGTH - 1; + let error = RequestStateCodec::try_new(vec![0; actual]).unwrap_err(); + + assert!( + matches!( + &error, + RequestStateError::KeyTooShort { + minimum: RequestStateCodec::MIN_KEY_LENGTH, + actual: error_actual, + } if *error_actual == actual + ), + "unexpected error: {error}" + ); + } + + #[test] + #[allow(deprecated)] + fn new_accepts_short_key_for_backward_compatibility() { + let codec = RequestStateCodec::new(b"key".to_vec()); + + assert_eq!(codec.open(&codec.seal(b"state")).unwrap(), b"state"); + } + + #[test] + fn seal_open_roundtrips_bytes() { + let codec = + RequestStateCodec::try_new(b"test-key-test-key-test-key-32byte".to_vec()).unwrap(); + let sealed = codec.seal(b"hello world"); + assert!(sealed.starts_with("rs1.")); + assert_eq!(codec.open(&sealed).unwrap(), b"hello world"); + } + + #[test] + fn seal_open_roundtrips_json() { + #[derive(serde::Serialize, serde::Deserialize, PartialEq, Debug)] + struct State { + tool: String, + round: u32, + } + let codec = + RequestStateCodec::try_new(b"another-strong-signing-key-here!!".to_vec()).unwrap(); + let state = State { + tool: "weather".into(), + round: 3, + }; + let sealed = codec.seal_json(&state).unwrap(); + let opened: State = codec.open_json(&sealed).unwrap(); + assert_eq!(opened, state); + } + + #[test] + fn empty_payload_roundtrips() { + let codec = + RequestStateCodec::try_new(b"empty-payload-test-signing-key!!".to_vec()).unwrap(); + let sealed = codec.seal(b""); + assert_eq!(codec.open(&sealed).unwrap(), b""); + } + + #[test] + fn tampered_payload_is_rejected() { + let codec = + RequestStateCodec::try_new(b"signing-key-signing-key-signing!!".to_vec()).unwrap(); + let sealed = codec.seal(b"amount=100"); + + // Replace the body section but keep the original tag. + let mut parts: Vec<&str> = sealed.split('.').collect(); + let forged_body = URL_SAFE_NO_PAD.encode(b"amount=999"); + parts[1] = &forged_body; + let forged = parts.join("."); + + assert!(matches!( + codec.open(&forged), + Err(RequestStateError::IntegrityCheckFailed) + )); + } + + #[test] + fn different_key_is_rejected() { + let signer = + RequestStateCodec::try_new(b"the-real-signing-key-value-here!!".to_vec()).unwrap(); + let attacker = + RequestStateCodec::try_new(b"a-totally-different-forged-key!!!".to_vec()).unwrap(); + let sealed = signer.seal(b"trusted"); + assert!(matches!( + attacker.open(&sealed), + Err(RequestStateError::IntegrityCheckFailed) + )); + } + + #[test] + fn appended_bytes_are_rejected() { + let codec = + RequestStateCodec::try_new(b"key-key-key-key-key-key-key-key!!".to_vec()).unwrap(); + let mut sealed = codec.seal(b"state"); + sealed.push('x'); + assert!(codec.open(&sealed).is_err()); + } + + #[test] + fn unsupported_version_is_malformed() { + let codec = + RequestStateCodec::try_new(b"wrong-version-test-signing-key!!".to_vec()).unwrap(); + let sealed = codec.seal(b"state"); + let bumped = sealed.replacen("rs1.", "rs3.", 1); + assert!(matches!( + codec.open(&bumped), + Err(RequestStateError::MalformedFormat) + )); + } + + #[test] + fn missing_sections_are_malformed() { + let codec = + RequestStateCodec::try_new(b"missing-sections-test-signing-key".to_vec()).unwrap(); + assert!(matches!( + codec.open("rs1"), + Err(RequestStateError::MalformedFormat) + )); + assert!(matches!( + codec.open("rs1.onlybody"), + Err(RequestStateError::MalformedFormat) + )); + assert!(matches!( + codec.open("rs1.a.b.c"), + Err(RequestStateError::MalformedFormat) + )); + } + + #[test] + fn non_base64_sections_are_invalid_encoding() { + let codec = + RequestStateCodec::try_new(b"invalid-base64-test-signing-key!!".to_vec()).unwrap(); + assert!(matches!( + codec.open("rs1.!!!!.!!!!"), + Err(RequestStateError::InvalidEncoding) + )); + } + + #[test] + fn debug_does_not_leak_key() { + let codec = + RequestStateCodec::try_new(b"super-secret-key-super-secret-key!!".to_vec()).unwrap(); + let rendered = format!("{codec:?}"); + assert!(!rendered.contains("super-secret-key")); + assert!(rendered.contains("redacted")); + } + + mod associated_data { + use super::*; + + #[test] + fn matching_context_opens() { + let codec = + RequestStateCodec::try_new(b"key-key-key-key-key-key-key-key!!".to_vec()).unwrap(); + let ctx = b"user:alice|tools/call:weather"; + let sealed = codec.seal_with(b"state", &SealOptions::new().associated_data(ctx)); + assert_eq!(codec.open_with(&sealed, ctx).unwrap(), b"state"); + } + + #[test] + fn different_context_is_rejected() { + let codec = + RequestStateCodec::try_new(b"key-key-key-key-key-key-key-key!!".to_vec()).unwrap(); + let sealed = + codec.seal_with(b"state", &SealOptions::new().associated_data(b"user:alice")); + assert!(matches!( + codec.open_with(&sealed, b"user:bob"), + Err(RequestStateError::IntegrityCheckFailed) + )); + } + + #[test] + fn missing_context_is_rejected() { + let codec = + RequestStateCodec::try_new(b"key-key-key-key-key-key-key-key!!".to_vec()).unwrap(); + let sealed = + codec.seal_with(b"state", &SealOptions::new().associated_data(b"user:alice")); + // Opening without the associated data must fail closed. + assert!(matches!( + codec.open(&sealed), + Err(RequestStateError::IntegrityCheckFailed) + )); + } + } + + mod ttl { + use super::*; + + const KEY: &[u8] = b"ttl-signing-key-ttl-signing-key!!"; + + #[test] + fn within_ttl_opens() { + let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); + let sealed = codec.seal_at( + b"state", + &SealOptions::new().ttl(Duration::from_secs(60)), + 1_000, + ); + // 30s later, still valid. + assert_eq!(codec.open_at(&sealed, &[], 31_000).unwrap(), b"state"); + } + + #[test] + fn past_ttl_is_expired() { + let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); + let sealed = codec.seal_at( + b"state", + &SealOptions::new().ttl(Duration::from_secs(60)), + 1_000, + ); + // 61s later, expired. + assert!(matches!( + codec.open_at(&sealed, &[], 62_000), + Err(RequestStateError::Expired) + )); + } + + #[test] + fn no_ttl_never_expires() { + let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); + let sealed = codec.seal_at(b"state", &SealOptions::new(), 1_000); + assert_eq!(codec.open_at(&sealed, &[], i64::MAX).unwrap(), b"state"); + } + + #[test] + fn ttl_and_associated_data_combine() { + let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); + let ctx = b"user:alice"; + let sealed = codec.seal_at( + b"state", + &SealOptions::new() + .associated_data(ctx) + .ttl(Duration::from_secs(60)), + 1_000, + ); + assert_eq!(codec.open_at(&sealed, ctx, 10_000).unwrap(), b"state"); + assert!(matches!( + codec.open_at(&sealed, b"user:bob", 10_000), + Err(RequestStateError::IntegrityCheckFailed) + )); + assert!(matches!( + codec.open_at(&sealed, ctx, 99_000), + Err(RequestStateError::Expired) + )); + } + } + + mod key_rotation { + use super::*; + + const KEY_A: &[u8] = b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + const KEY_B: &[u8] = b"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"; + const KEY_C: &[u8] = b"cccccccccccccccccccccccccccccccc"; + const KAT_KEY: &[u8] = b"0123456789abcdef0123456789abcdef"; + const KAT_KID: &str = "rotation-key-2026-08"; + const KAT_AD: &[u8] = b"user:alice|request:weather"; + const KAT_NOW_MS: i64 = 1_700_000_000_000; + + fn two_key_ring(active: &str) -> RequestStateCodec { + RequestStateCodec::new_with_keyring(active, [("a", KEY_A), ("b", KEY_B)]).unwrap() + } + + fn body(expiry: i64, payload: &[u8]) -> Vec { + let mut body = expiry.to_be_bytes().to_vec(); + body.extend_from_slice(payload); + body + } + + fn replace_segment(token: &str, index: usize, replacement: &str) -> String { + let mut parts: Vec = token.split('.').map(str::to_owned).collect(); + parts[index] = replacement.to_owned(); + parts.join(".") + } + + fn test_mac( + key: &[u8], + domain: &[u8], + kid: Option<&[u8]>, + associated_data: &[u8], + body: &[u8], + ) -> Vec { + let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts keys of any length"); + mac.update(domain); + if let Some(kid) = kid { + mac.update(&(kid.len() as u64).to_be_bytes()); + mac.update(kid); + } + mac.update(&(associated_data.len() as u64).to_be_bytes()); + mac.update(associated_data); + mac.update(body); + mac.finalize().into_bytes().to_vec() + } + + fn raw_rs1(body: &[u8], tag: &[u8]) -> String { + format!( + "rs1.{}.{}", + URL_SAFE_NO_PAD.encode(body), + URL_SAFE_NO_PAD.encode(tag) + ) + } + + fn raw_rs2(kid: &[u8], body: &[u8], tag: &[u8]) -> String { + format!( + "rs2.{}.{}.{}", + URL_SAFE_NO_PAD.encode(kid), + URL_SAFE_NO_PAD.encode(body), + URL_SAFE_NO_PAD.encode(tag) + ) + } + + #[test] + fn rs1_known_answer_is_unchanged() { + // Independently checked with Python stdlib HMAC and Ruby OpenSSL. + let codec = RequestStateCodec::try_new(KAT_KEY).unwrap(); + let sealed = codec.seal_at( + b"step=2", + &SealOptions::new() + .associated_data(KAT_AD) + .ttl(Duration::from_secs(90)), + KAT_NOW_MS, + ); + assert_eq!( + sealed, + "rs1.AAABi8_mx5BzdGVwPTI.GQgS0X7mtSz8ZOy_kld2Zjuc4gAMGpBL74EghWI36IQ" + ); + } + + #[test] + fn rs2_known_answer_matches_wire_specification() { + // Independently checked with Python stdlib HMAC and Ruby OpenSSL. + let codec = RequestStateCodec::new_with_keyring(KAT_KID, [(KAT_KID, KAT_KEY)]).unwrap(); + let sealed = codec.seal_at( + b"step=2", + &SealOptions::new() + .associated_data(KAT_AD) + .ttl(Duration::from_secs(90)), + KAT_NOW_MS, + ); + assert_eq!( + sealed, + "rs2.cm90YXRpb24ta2V5LTIwMjYtMDg.AAABi8_mx5BzdGVwPTI.twv2acu7lKXqebyrmit-JrHHZm-BkKlBQEMMtL3lHk8" + ); + } + + #[test] + fn rs2_roundtrips_with_associated_data_and_ttl() { + let codec = two_key_ring("a"); + let options = SealOptions::new() + .associated_data(b"user:alice") + .ttl(Duration::from_secs(60)); + + let sealed = codec.seal_at(b"state", &options, 1_000); + assert!(sealed.starts_with("rs2.YQ.")); + assert_eq!( + codec.open_at(&sealed, b"user:alice", 30_000).unwrap(), + b"state" + ); + assert!(matches!( + codec.open_at(&sealed, b"user:bob", 30_000), + Err(RequestStateError::IntegrityCheckFailed) + )); + assert!(matches!( + codec.open_at(&sealed, b"user:alice", 70_000), + Err(RequestStateError::Expired) + )); + } + + #[test] + fn rotation_keeps_previous_rs2_key_verifiable() { + let signer_a = two_key_ring("a"); + let token_a = signer_a.seal(b"state-a"); + + let signer_b = two_key_ring("b"); + assert_eq!(signer_b.open(&token_a).unwrap(), b"state-a"); + let token_b = signer_b.seal(b"state-b"); + assert!(token_b.starts_with("rs2.Yg.")); + assert_eq!(signer_b.open(&token_b).unwrap(), b"state-b"); + + let without_a = RequestStateCodec::new_with_keyring("b", [("b", KEY_B)]).unwrap(); + assert!(matches!( + without_a.open(&token_a), + Err(RequestStateError::UnknownKeyId) + )); + } + + #[test] + fn rolling_migration_is_bidirectionally_compatible() { + let old = RequestStateCodec::try_new(KEY_A).unwrap(); + let transitional = + RequestStateCodec::new_with_keyring("new", [("old", KEY_A), ("new", KEY_B)]) + .unwrap() + .with_rs1_signing("old") + .unwrap(); + let promoted = + RequestStateCodec::new_with_keyring("new", [("old", KEY_A), ("new", KEY_B)]) + .unwrap() + .with_rs1_fallback("old") + .unwrap(); + let retired = RequestStateCodec::new_with_keyring("new", [("new", KEY_B)]).unwrap(); + + let old_token = old.seal(b"old"); + assert_eq!(transitional.open(&old_token).unwrap(), b"old"); + + let transitional_token = transitional.seal(b"transition"); + assert!(transitional_token.starts_with("rs1.")); + assert_eq!(old.open(&transitional_token).unwrap(), b"transition"); + assert_eq!(promoted.open(&transitional_token).unwrap(), b"transition"); + + let promoted_token = promoted.seal(b"promoted"); + assert!(promoted_token.starts_with("rs2.")); + assert_eq!(transitional.open(&promoted_token).unwrap(), b"promoted"); + assert_eq!(retired.open(&promoted_token).unwrap(), b"promoted"); + assert!(matches!( + retired.open(&old_token), + Err(RequestStateError::UnknownKeyId) + )); + } + + #[test] + fn multiple_legacy_fallbacks_are_supported() { + let legacy_a = RequestStateCodec::try_new(KEY_A).unwrap().seal(b"a"); + let legacy_b = RequestStateCodec::try_new(KEY_B).unwrap().seal(b"b"); + let legacy_c = RequestStateCodec::try_new(KEY_C).unwrap().seal(b"c"); + let ring = RequestStateCodec::new_with_keyring( + "new", + [("a", KEY_A), ("b", KEY_B), ("new", KEY_C)], + ) + .unwrap() + .with_rs1_fallback("a") + .unwrap() + .with_rs1_fallback("b") + .unwrap(); + + assert_eq!(ring.open(&legacy_a).unwrap(), b"a"); + assert_eq!(ring.open(&legacy_b).unwrap(), b"b"); + assert!(matches!( + ring.open(&legacy_c), + Err(RequestStateError::IntegrityCheckFailed) + )); + } + + #[test] + fn kid_is_authenticated_even_when_ids_share_key_bytes() { + let codec = + RequestStateCodec::new_with_keyring("a", [("a", KEY_A), ("b", KEY_A)]).unwrap(); + let sealed = codec.seal(b"state"); + let swapped = replace_segment(&sealed, 1, &URL_SAFE_NO_PAD.encode(b"b")); + assert!(matches!( + codec.open(&swapped), + Err(RequestStateError::IntegrityCheckFailed) + )); + } + + #[test] + fn independent_rs2_segment_tampering_is_rejected() { + let codec = two_key_ring("a"); + let sealed = codec.seal(b"state"); + + let swapped_kid = replace_segment(&sealed, 1, &URL_SAFE_NO_PAD.encode(b"b")); + let tampered_body = replace_segment(&sealed, 2, &URL_SAFE_NO_PAD.encode(b"changed")); + let tampered_tag = replace_segment(&sealed, 3, &URL_SAFE_NO_PAD.encode([0_u8; 32])); + + for tampered in [swapped_kid, tampered_body, tampered_tag] { + assert!(matches!( + codec.open(&tampered), + Err(RequestStateError::IntegrityCheckFailed) + )); + } + } + + #[test] + fn version_domains_are_cryptographically_separate() { + let token_body = body(0, b"state"); + + let wrong_v2_tag = test_mac(KEY_A, DOMAIN_V1, Some(b"a"), b"", &token_body); + let wrong_v2 = raw_rs2(b"a", &token_body, &wrong_v2_tag); + assert!(matches!( + two_key_ring("a").open(&wrong_v2), + Err(RequestStateError::IntegrityCheckFailed) + )); + + let wrong_v1_tag = test_mac(KEY_A, DOMAIN_V2, None, b"", &token_body); + let wrong_v1 = raw_rs1(&token_body, &wrong_v1_tag); + assert!(matches!( + RequestStateCodec::try_new(KEY_A).unwrap().open(&wrong_v1), + Err(RequestStateError::IntegrityCheckFailed) + )); + } + + #[test] + fn constructor_invariants_are_enforced() { + let empty = + RequestStateCodec::new_with_keyring("a", std::iter::empty::<(&str, &[u8])>()); + assert!(matches!(empty, Err(RequestStateError::InvalidKeyring(_)))); + + let duplicate = RequestStateCodec::new_with_keyring("a", [("a", KEY_A), ("a", KEY_B)]); + assert!(matches!( + duplicate, + Err(RequestStateError::InvalidKeyring(_)) + )); + + let empty_kid = RequestStateCodec::new_with_keyring("", [("", KEY_A)]); + assert!(matches!( + empty_kid, + Err(RequestStateError::InvalidKeyring(_)) + )); + + let oversized = "x".repeat(MAX_KID_LEN + 1); + let oversized_kid = + RequestStateCodec::new_with_keyring(oversized.clone(), [(oversized, KEY_A)]); + assert!(matches!( + oversized_kid, + Err(RequestStateError::InvalidKeyring(_)) + )); + + let absent_active = RequestStateCodec::new_with_keyring("missing", [("a", KEY_A)]); + assert!(matches!( + absent_active, + Err(RequestStateError::InvalidKeyring(_)) + )); + + let short_key = RequestStateCodec::new_with_keyring( + "a", + [("a", vec![0; RequestStateCodec::MIN_KEY_LENGTH - 1])], + ); + assert!(matches!( + short_key, + Err(RequestStateError::KeyTooShort { .. }) + )); + + assert!(matches!( + RequestStateCodec::try_new(KEY_A) + .unwrap() + .with_rs1_signing("a"), + Err(RequestStateError::InvalidKeyring(_)) + )); + assert!(matches!( + two_key_ring("a").with_rs1_signing("missing"), + Err(RequestStateError::InvalidKeyring(_)) + )); + assert!(matches!( + RequestStateCodec::try_new(KEY_A) + .unwrap() + .with_rs1_fallback("a"), + Err(RequestStateError::InvalidKeyring(_)) + )); + assert!(matches!( + two_key_ring("a").with_rs1_fallback("missing"), + Err(RequestStateError::InvalidKeyring(_)) + )); + } + + #[test] + fn maximum_length_kid_roundtrips_and_oversized_wire_kid_is_rejected() { + let max_kid = "x".repeat(MAX_KID_LEN); + let codec = + RequestStateCodec::new_with_keyring(max_kid.clone(), [(max_kid, KEY_A)]).unwrap(); + let token = codec.seal(b"state"); + assert_eq!(codec.open(&token).unwrap(), b"state"); + + let encoded_oversized = URL_SAFE_NO_PAD.encode("x".repeat(MAX_KID_LEN + 1)); + assert!(encoded_oversized.len() > MAX_ENCODED_KID_LEN); + let oversized = format!("rs2.{encoded_oversized}.!!!!.!!!!"); + assert!(matches!( + codec.open(&oversized), + Err(RequestStateError::InvalidKeyId) + )); + + let overlong_invalid_base64 = + format!("rs2.{}.!!!!.!!!!", "!".repeat(MAX_ENCODED_KID_LEN + 1)); + assert!(matches!( + codec.open(&overlong_invalid_base64), + Err(RequestStateError::InvalidKeyId) + )); + } + + #[test] + fn parser_rejects_invalid_tokens_and_unavailable_keys() { + let codec = two_key_ring("a"); + for malformed in [ + "rs2", + "rs2.YQ", + "rs2.YQ.body", + "rs2.YQ.body.tag.extra", + "rs3.YQ.body.tag", + ] { + assert!(matches!( + codec.open(malformed), + Err(RequestStateError::MalformedFormat) + )); + } + + let valid_rs2 = codec.seal(b"state"); + assert!(matches!( + codec.open(&replace_segment(&valid_rs2, 1, "!!!!")), + Err(RequestStateError::InvalidEncoding) + )); + assert!(matches!( + codec.open(&replace_segment(&valid_rs2, 1, "")), + Err(RequestStateError::InvalidKeyId) + )); + let non_utf8 = URL_SAFE_NO_PAD.encode([0xff]); + assert!(matches!( + codec.open(&replace_segment(&valid_rs2, 1, &non_utf8)), + Err(RequestStateError::InvalidKeyId) + )); + + assert!(matches!( + RequestStateCodec::try_new(KEY_A).unwrap().open(&valid_rs2), + Err(RequestStateError::UnknownKeyId) + )); + + let valid_rs1 = RequestStateCodec::try_new(KEY_A).unwrap().seal(b"state"); + assert!(matches!( + codec.open(&valid_rs1), + Err(RequestStateError::UnknownKeyId) + )); + } + + #[test] + fn authenticated_short_body_is_malformed() { + let short_body = b"short"; + let tag = RequestStateCodec::mac_v2(KEY_A, b"a", b"", short_body) + .finalize() + .into_bytes(); + let token = raw_rs2(b"a", short_body, tag.as_slice()); + assert!(matches!( + two_key_ring("a").open(&token), + Err(RequestStateError::MalformedFormat) + )); + } + + #[test] + fn serde_is_reached_only_after_integrity_verification() { + let codec = two_key_ring("a"); + let invalid_json = codec.seal(b"{not-json"); + let opened: Result = codec.open_json(&invalid_json); + assert!(matches!(opened, Err(RequestStateError::Deserialization(_)))); + + let valid_json = codec.seal(b"{}"); + let tampered_body = body(0, b"{not-json"); + let tampered = replace_segment(&valid_json, 2, &URL_SAFE_NO_PAD.encode(tampered_body)); + let opened: Result = codec.open_json(&tampered); + assert!(matches!( + opened, + Err(RequestStateError::IntegrityCheckFailed) + )); + } + + #[test] + fn ring_debug_redacts_all_key_material() { + let codec = two_key_ring("b").with_rs1_fallback("a").unwrap(); + let rendered = format!("{codec:?}"); + assert!(!rendered.contains(std::str::from_utf8(KEY_A).unwrap())); + assert!(!rendered.contains(std::str::from_utf8(KEY_B).unwrap())); + assert!(rendered.contains("redacted")); + } + + #[test] + fn mutated_tokens_do_not_panic() { + let codec = two_key_ring("a").with_rs1_fallback("a").unwrap(); + let valid = [ + RequestStateCodec::try_new(KEY_A).unwrap().seal(b"state"), + codec.seal(b"state"), + ]; + let mut corpus = vec![ + String::new(), + ".".to_owned(), + "...".to_owned(), + "rs1".to_owned(), + "rs2".to_owned(), + "💥".to_owned(), + "rs2.💥...".to_owned(), + ]; + + for token in valid { + for end in 0..=token.len() { + corpus.push(token[..end].to_owned()); + } + for index in 0..token.len() { + let mut bytes = token.as_bytes().to_vec(); + bytes[index] = b'!'; + corpus.push(String::from_utf8(bytes).expect("token is ASCII")); + } + } + + for candidate in corpus { + let _ = codec.open(&candidate); + let _ = codec.open_with(&candidate, b"context"); + } + } + } +} diff --git a/crates/rmcp/src/model/resource.rs b/crates/rmcp/src/model/resource.rs index a5ad950..9c27dce 100644 --- a/crates/rmcp/src/model/resource.rs +++ b/crates/rmcp/src/model/resource.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -use super::{Annotations, Icon, Meta}; +use super::{Annotations, Icon, MetaObject}; /// A known resource that the server is capable of reading (spec `Resource`). /// @@ -31,7 +31,7 @@ pub struct Resource { pub icons: Option>, /// Optional protocol-level metadata for this resource. #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// Optional annotations describing how the client should use this resource. #[serde(skip_serializing_if = "Option::is_none")] pub annotations: Option, @@ -77,7 +77,7 @@ impl Resource { self } - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } @@ -112,7 +112,7 @@ pub struct ResourceTemplate { pub icons: Option>, /// Optional protocol-level metadata for this resource template. #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, /// Optional annotations describing how the client should use this template. #[serde(skip_serializing_if = "Option::is_none")] pub annotations: Option, @@ -152,7 +152,7 @@ impl ResourceTemplate { self } - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } @@ -176,7 +176,7 @@ pub enum ResourceContents { mime_type: Option, text: String, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - meta: Option, + meta: Option, }, #[serde(rename_all = "camelCase")] BlobResourceContents { @@ -185,7 +185,7 @@ pub enum ResourceContents { mime_type: Option, blob: String, #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - meta: Option, + meta: Option, }, } @@ -193,7 +193,7 @@ impl ResourceContents { pub fn text(text: impl Into, uri: impl Into) -> Self { Self::TextResourceContents { uri: uri.into(), - mime_type: Some("text".into()), + mime_type: Some("text/plain".into()), text: text.into(), meta: None, } @@ -216,7 +216,7 @@ impl ResourceContents { self } - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { match &mut self { Self::TextResourceContents { meta: m, .. } => *m = Some(meta), Self::BlobResourceContents { meta: m, .. } => *m = Some(meta), @@ -303,7 +303,7 @@ mod tests { #[test] fn test_resource_template_with_meta() { let resource_template = - ResourceTemplate::new("file:///{path}", "template").with_meta(Meta::default()); + ResourceTemplate::new("file:///{path}", "template").with_meta(MetaObject::default()); let json = serde_json::to_value(&resource_template).unwrap(); assert!(json.get("_meta").is_some()); } diff --git a/crates/rmcp/src/model/serde_impl.rs b/crates/rmcp/src/model/serde_impl.rs index f8996f3..bfd3b57 100644 --- a/crates/rmcp/src/model/serde_impl.rs +++ b/crates/rmcp/src/model/serde_impl.rs @@ -3,12 +3,19 @@ use std::borrow::Cow; use serde::{Deserialize, Serialize}; use super::{ - CustomNotification, CustomRequest, Extensions, Meta, Notification, NotificationNoParam, - Request, RequestNoParam, RequestOptionalParam, + CustomNotification, CustomRequest, Extensions, JsonObject, MetaObject, Notification, + NotificationMetaObject, NotificationNoParam, Request, RequestMetaObject, RequestNoParam, + RequestOptionalParam, }; + +/// Wire-side view of `params`: the `_meta` map plus the remaining fields. +/// +/// All metadata types are transparent wrappers over [`JsonObject`], so the +/// serde plumbing works on the raw map; call sites wrap/unwrap the typed +/// metadata ([`RequestMetaObject`] / [`NotificationMetaObject`]). #[derive(Deserialize)] struct WithMeta<'a, P> { - _meta: Option>, + _meta: Option>, #[serde(flatten)] _rest: P, } @@ -25,7 +32,7 @@ impl Serialize for WithMeta<'_, P> { serde_json::to_value(&self._rest).map_err(serde::ser::Error::custom)?; // Extract _meta from the serialized params (if it's an object containing one) - let params_meta: Option = rest_value + let params_meta: Option = rest_value .as_object_mut() .and_then(|obj| obj.remove("_meta")) .and_then(|v| serde_json::from_value(v).ok()); @@ -78,6 +85,56 @@ struct ProxyNoParam { method: M, } +/// Combine the message-specific `_meta` map with a legacy [`MetaObject`] +/// extension so metadata stored in [`Extensions`] is not lost on the wire. +/// On key conflicts the message-specific map wins. +fn merge_legacy_meta<'a>( + typed: Option<&'a JsonObject>, + extensions: &'a Extensions, +) -> Option> { + let legacy = extensions.get::().map(|meta| &meta.0); + match (typed, legacy) { + (Some(typed), None) => Some(Cow::Borrowed(typed)), + (None, Some(legacy)) => Some(Cow::Borrowed(legacy)), + (Some(typed), Some(legacy)) => { + let mut merged = legacy.clone(); + merged.extend( + typed + .iter() + .map(|(key, value)| (key.clone(), value.clone())), + ); + Some(Cow::Owned(merged)) + } + (None, None) => None, + } +} + +/// Borrow the request `_meta` map from extensions, if any. +fn request_meta(extensions: &Extensions) -> Option> { + let typed = extensions.get::().map(|meta| &meta.0.0); + merge_legacy_meta(typed, extensions) +} + +/// Borrow the notification `_meta` map from extensions, if any. +fn notification_meta(extensions: &Extensions) -> Option> { + let typed = extensions + .get::() + .map(|meta| &meta.0.0); + merge_legacy_meta(typed, extensions) +} + +/// Build extensions holding a typed metadata map deserialized from `params._meta`. +fn extensions_with_meta(meta: Option>) -> Extensions +where + T: From + Clone + Send + Sync + 'static, +{ + let mut extensions = Extensions::new(); + if let Some(meta) = meta { + extensions.insert(T::from(meta.into_owned())); + } + extensions +} + impl Serialize for Request where M: Serialize, @@ -87,14 +144,12 @@ where where S: serde::Serializer, { - let extensions = &self.extensions; - let _meta = extensions.get::().map(Cow::Borrowed); Proxy::serialize( &Proxy { method: &self.method, params: WithMeta { _rest: &self.params, - _meta, + _meta: request_meta(&self.extensions), }, }, serializer, @@ -112,13 +167,8 @@ where D: serde::Deserializer<'de>, { let body = Proxy::deserialize(deserializer)?; - let _meta = body.params._meta.map(|m| m.into_owned()); - let mut extensions = Extensions::new(); - if let Some(meta) = _meta { - extensions.insert(meta); - } Ok(Request { - extensions, + extensions: extensions_with_meta::(body.params._meta), method: body.method, params: body.params._rest, }) @@ -134,14 +184,12 @@ where where S: serde::Serializer, { - let extensions = &self.extensions; - let _meta = extensions.get::().map(Cow::Borrowed); Proxy::serialize( &Proxy { method: &self.method, params: WithMeta { _rest: &self.params, - _meta, + _meta: request_meta(&self.extensions), }, }, serializer, @@ -163,14 +211,10 @@ where let mut _meta = None; if let Some(body_params) = body.params { params = body_params._rest; - _meta = body_params._meta.map(|m| m.into_owned()); - } - let mut extensions = Extensions::new(); - if let Some(meta) = _meta { - extensions.insert(meta); + _meta = body_params._meta; } Ok(RequestOptionalParam { - extensions, + extensions: extensions_with_meta::(_meta), method: body.method, params, }) @@ -185,14 +229,26 @@ where where S: serde::Serializer, { - let extensions = &self.extensions; - let _meta = extensions.get::().map(Cow::Borrowed); - ProxyNoParam::serialize( - &ProxyNoParam { - method: &self.method, - }, - serializer, - ) + // Emit `params` only when metadata is present, so the wire shape of + // meta-less requests stays `{"method": ...}`. + match request_meta(&self.extensions) { + Some(_meta) => Proxy::serialize( + &Proxy { + method: &self.method, + params: WithMeta { + _meta: Some(_meta), + _rest: (), + }, + }, + serializer, + ), + None => ProxyNoParam::serialize( + &ProxyNoParam { + method: &self.method, + }, + serializer, + ), + } } } @@ -204,10 +260,10 @@ where where D: serde::Deserializer<'de>, { - let body = ProxyNoParam::<_>::deserialize(deserializer)?; - let extensions = Extensions::new(); + let body = ProxyOptionalParam::<'_, _, Option>::deserialize(deserializer)?; + let _meta = body.params.and_then(|params| params._meta); Ok(RequestNoParam { - extensions, + extensions: extensions_with_meta::(_meta), method: body.method, }) } @@ -222,14 +278,12 @@ where where S: serde::Serializer, { - let extensions = &self.extensions; - let _meta = extensions.get::().map(Cow::Borrowed); Proxy::serialize( &Proxy { method: &self.method, params: WithMeta { _rest: &self.params, - _meta, + _meta: notification_meta(&self.extensions), }, }, serializer, @@ -248,10 +302,7 @@ where { let body = ProxyOptionalParam::<'_, _, R>::deserialize(deserializer)?; let (_meta, params) = match body.params { - Some(with_meta) => { - let meta = with_meta._meta.map(|m| m.into_owned()); - (meta, with_meta._rest) - } + Some(with_meta) => (with_meta._meta, with_meta._rest), None => { // JSON-RPC 2.0: params is optional. Treat absent params as {}. let empty = serde_json::Value::Object(serde_json::Map::new()); @@ -259,12 +310,8 @@ where (None, r) } }; - let mut extensions = Extensions::new(); - if let Some(meta) = _meta { - extensions.insert(meta); - } Ok(Notification { - extensions, + extensions: extensions_with_meta::(_meta), method: body.method, params, }) @@ -279,14 +326,26 @@ where where S: serde::Serializer, { - let extensions = &self.extensions; - let _meta = extensions.get::().map(Cow::Borrowed); - ProxyNoParam::serialize( - &ProxyNoParam { - method: &self.method, - }, - serializer, - ) + // Emit `params` only when metadata is present, so the wire shape of + // meta-less notifications stays `{"method": ...}`. + match notification_meta(&self.extensions) { + Some(_meta) => Proxy::serialize( + &Proxy { + method: &self.method, + params: WithMeta { + _meta: Some(_meta), + _rest: (), + }, + }, + serializer, + ), + None => ProxyNoParam::serialize( + &ProxyNoParam { + method: &self.method, + }, + serializer, + ), + } } } @@ -298,10 +357,10 @@ where where D: serde::Deserializer<'de>, { - let body = ProxyNoParam::<_>::deserialize(deserializer)?; - let extensions = Extensions::new(); + let body = ProxyOptionalParam::<'_, _, Option>::deserialize(deserializer)?; + let _meta = body.params.and_then(|params| params._meta); Ok(NotificationNoParam { - extensions, + extensions: extensions_with_meta::(_meta), method: body.method, }) } @@ -312,8 +371,7 @@ impl Serialize for CustomRequest { where S: serde::Serializer, { - let extensions = &self.extensions; - let _meta = extensions.get::().map(Cow::Borrowed); + let _meta = request_meta(&self.extensions); let params = self.params.as_ref(); let params = if _meta.is_some() || params.is_some() { @@ -346,14 +404,10 @@ impl<'de> Deserialize<'de> for CustomRequest { let mut _meta = None; if let Some(body_params) = body.params { params = body_params._rest; - _meta = body_params._meta.map(|m| m.into_owned()); - } - let mut extensions = Extensions::new(); - if let Some(meta) = _meta { - extensions.insert(meta); + _meta = body_params._meta; } Ok(CustomRequest { - extensions, + extensions: extensions_with_meta::(_meta), method: body.method, params, }) @@ -365,8 +419,7 @@ impl Serialize for CustomNotification { where S: serde::Serializer, { - let extensions = &self.extensions; - let _meta = extensions.get::().map(Cow::Borrowed); + let _meta = notification_meta(&self.extensions); let params = self.params.as_ref(); let params = if _meta.is_some() || params.is_some() { @@ -399,14 +452,10 @@ impl<'de> Deserialize<'de> for CustomNotification { let mut _meta = None; if let Some(body_params) = body.params { params = body_params._rest; - _meta = body_params._meta.map(|m| m.into_owned()); - } - let mut extensions = Extensions::new(); - if let Some(meta) = _meta { - extensions.insert(meta); + _meta = body_params._meta; } Ok(CustomNotification { - extensions, + extensions: extensions_with_meta::(_meta), method: body.method, params, }) @@ -418,7 +467,8 @@ mod test { use serde_json::json; use crate::model::{ - CallToolRequest, CallToolRequestParams, CustomRequest, Extensions, ListToolsRequest, Meta, + CallToolRequest, CallToolRequestParams, CustomRequest, Extensions, InitializedNotification, + ListToolsRequest, NotificationMetaObject, PingRequest, RequestMetaObject, }; #[test] @@ -436,12 +486,12 @@ mod test { // When both extensions and params contain _meta, the output should have // a single merged _meta key (not two separate ones). let mut extensions = Extensions::new(); - let mut ext_meta = Meta::new(); - ext_meta.0.insert("traceId".to_string(), json!("abc")); + let mut ext_meta = RequestMetaObject::new(); + ext_meta.insert("traceId".to_string(), json!("abc")); extensions.insert(ext_meta); - let mut params_meta = Meta::new(); - params_meta.0.insert("progressToken".to_string(), json!(1)); + let mut params_meta = RequestMetaObject::new(); + params_meta.insert("progressToken".to_string(), json!(1)); let req = CallToolRequest { extensions, @@ -450,7 +500,8 @@ mod test { meta: Some(params_meta), name: "my_tool".into(), arguments: None, - task: None, + input_responses: None, + request_state: None, }, }; @@ -477,8 +528,8 @@ mod test { #[test] fn test_meta_only_from_extensions() { let mut extensions = Extensions::new(); - let mut ext_meta = Meta::new(); - ext_meta.0.insert("traceId".to_string(), json!("ext-only")); + let mut ext_meta = RequestMetaObject::new(); + ext_meta.insert("traceId".to_string(), json!("ext-only")); extensions.insert(ext_meta); let req = CallToolRequest { @@ -488,7 +539,8 @@ mod test { meta: None, name: "my_tool".into(), arguments: None, - task: None, + input_responses: None, + request_state: None, }, }; @@ -499,8 +551,8 @@ mod test { #[test] fn test_meta_only_from_params() { - let mut params_meta = Meta::new(); - params_meta.0.insert("progressToken".to_string(), json!(42)); + let mut params_meta = RequestMetaObject::new(); + params_meta.insert("progressToken".to_string(), json!(42)); let req = CallToolRequest { extensions: Extensions::new(), @@ -509,7 +561,8 @@ mod test { meta: Some(params_meta), name: "my_tool".into(), arguments: None, - task: None, + input_responses: None, + request_state: None, }, }; @@ -527,7 +580,8 @@ mod test { meta: None, name: "my_tool".into(), arguments: None, - task: None, + input_responses: None, + request_state: None, }, }; @@ -542,19 +596,13 @@ mod test { fn test_extensions_meta_takes_priority_on_conflict() { // When both sources have the same key, extensions should win. let mut extensions = Extensions::new(); - let mut ext_meta = Meta::new(); - ext_meta - .0 - .insert("shared_key".to_string(), json!("from_extensions")); + let mut ext_meta = RequestMetaObject::new(); + ext_meta.insert("shared_key".to_string(), json!("from_extensions")); extensions.insert(ext_meta); - let mut params_meta = Meta::new(); - params_meta - .0 - .insert("shared_key".to_string(), json!("from_params")); - params_meta - .0 - .insert("params_only".to_string(), json!("kept")); + let mut params_meta = RequestMetaObject::new(); + params_meta.insert("shared_key".to_string(), json!("from_params")); + params_meta.insert("params_only".to_string(), json!("kept")); let req = CallToolRequest { extensions, @@ -563,7 +611,8 @@ mod test { meta: Some(params_meta), name: "my_tool".into(), arguments: None, - task: None, + input_responses: None, + request_state: None, }, }; @@ -576,10 +625,8 @@ mod test { #[test] fn test_round_trip_preserves_meta() { let mut extensions = Extensions::new(); - let mut ext_meta = Meta::new(); - ext_meta - .0 - .insert("traceId".to_string(), json!("round-trip")); + let mut ext_meta = RequestMetaObject::new(); + ext_meta.insert("traceId".to_string(), json!("round-trip")); extensions.insert(ext_meta); let req = CallToolRequest { @@ -589,7 +636,8 @@ mod test { meta: None, name: "my_tool".into(), arguments: Some(serde_json::Map::from_iter([("x".to_string(), json!(1))])), - task: None, + input_responses: None, + request_state: None, }, }; @@ -597,8 +645,8 @@ mod test { let deserialized: CallToolRequest = serde_json::from_str(&serialized).unwrap(); // Extensions should have the meta after round-trip - let meta = deserialized.extensions.get::().unwrap(); - assert_eq!(meta.0.get("traceId").unwrap(), "round-trip"); + let meta = deserialized.extensions.get::().unwrap(); + assert_eq!(meta.get("traceId").unwrap(), "round-trip"); // Params should be preserved assert_eq!(deserialized.params.name, "my_tool"); @@ -618,10 +666,8 @@ mod test { fn test_custom_request_no_duplicate_meta() { // CustomRequest uses Option as params — verify no duplicate _meta. let mut extensions = Extensions::new(); - let mut ext_meta = Meta::new(); - ext_meta - .0 - .insert("traceId".to_string(), json!("custom-ext")); + let mut ext_meta = RequestMetaObject::new(); + ext_meta.insert("traceId".to_string(), json!("custom-ext")); extensions.insert(ext_meta); let params = Some(json!({ @@ -648,4 +694,158 @@ mod test { assert_eq!(meta.get("traceId").unwrap(), "custom-ext"); assert_eq!(meta.get("progressToken").unwrap(), 99); } + + #[test] + fn test_request_no_param_meta_round_trip() { + // Ping-shaped requests must carry `params._meta` on the wire. + let mut extensions = Extensions::new(); + let mut meta = RequestMetaObject::new(); + meta.insert("traceId".to_string(), json!("ping-trace")); + extensions.insert(meta); + + let req = PingRequest { + method: Default::default(), + extensions, + }; + + let value = serde_json::to_value(&req).unwrap(); + assert_eq!(value["params"]["_meta"]["traceId"], json!("ping-trace")); + + let deserialized: PingRequest = serde_json::from_value(value).unwrap(); + let meta = deserialized + .extensions + .get::() + .expect("meta should survive the round-trip"); + assert_eq!(meta.get("traceId").unwrap(), &json!("ping-trace")); + } + + #[test] + fn test_request_no_param_without_meta_has_no_params_key() { + let req = PingRequest { + method: Default::default(), + extensions: Extensions::new(), + }; + let value = serde_json::to_value(&req).unwrap(); + assert!( + value.get("params").is_none(), + "meta-less no-param requests must keep the historical wire shape: {value}" + ); + } + + #[test] + fn test_notification_no_param_meta_round_trip() { + // Initialized-shaped notifications must carry `params._meta` on the wire. + let mut extensions = Extensions::new(); + let mut meta = NotificationMetaObject::new(); + meta.insert("traceId".to_string(), json!("init-trace")); + extensions.insert(meta); + + let notification = InitializedNotification { + method: Default::default(), + extensions, + }; + + let value = serde_json::to_value(¬ification).unwrap(); + assert_eq!(value["params"]["_meta"]["traceId"], json!("init-trace")); + + let deserialized: InitializedNotification = serde_json::from_value(value).unwrap(); + let meta = deserialized + .extensions + .get::() + .expect("meta should survive the round-trip"); + assert_eq!(meta.get("traceId").unwrap(), &json!("init-trace")); + } + + #[test] + fn test_notification_no_param_without_meta_has_no_params_key() { + let notification = InitializedNotification { + method: Default::default(), + extensions: Extensions::new(), + }; + let value = serde_json::to_value(¬ification).unwrap(); + assert!( + value.get("params").is_none(), + "meta-less no-param notifications must keep the historical wire shape: {value}" + ); + } + + #[test] + fn test_no_param_ignores_unknown_params_fields() { + // Old/foreign peers may send params without _meta; both shapes must parse. + let _req: PingRequest = + serde_json::from_value(json!({"method": "ping", "params": {}})).unwrap(); + let _req: PingRequest = + serde_json::from_value(json!({"method": "ping", "params": {"unknown": 1}})).unwrap(); + let _req: PingRequest = serde_json::from_value(json!({"method": "ping"})).unwrap(); + } + + #[test] + fn test_legacy_meta_extension_still_serializes() { + let mut extensions = Extensions::new(); + let mut legacy = crate::model::MetaObject::new(); + legacy.insert("traceId".to_string(), json!("legacy")); + extensions.insert(legacy); + + let req = CallToolRequest { + extensions, + method: Default::default(), + params: CallToolRequestParams { + meta: None, + name: "my_tool".into(), + arguments: None, + input_responses: None, + request_state: None, + }, + }; + + let value = serde_json::to_value(&req).unwrap(); + assert_eq!(value["params"]["_meta"]["traceId"], json!("legacy")); + } + + #[test] + fn test_typed_meta_wins_over_legacy_extension_on_conflict() { + let mut extensions = Extensions::new(); + let mut legacy = crate::model::MetaObject::new(); + legacy.insert("shared".to_string(), json!("legacy")); + legacy.insert("legacy_only".to_string(), json!("kept")); + extensions.insert(legacy); + let mut typed = RequestMetaObject::new(); + typed.insert("shared".to_string(), json!("typed")); + extensions.insert(typed); + + let req = CallToolRequest { + extensions, + method: Default::default(), + params: CallToolRequestParams { + meta: None, + name: "my_tool".into(), + arguments: None, + input_responses: None, + request_state: None, + }, + }; + + let value = serde_json::to_value(&req).unwrap(); + let meta = value["params"]["_meta"].as_object().unwrap(); + assert_eq!(meta.get("shared").unwrap(), "typed"); + assert_eq!(meta.get("legacy_only").unwrap(), "kept"); + } + + #[test] + fn test_arbitrary_meta_keys_round_trip_unchanged() { + let input = json!({ + "method": "tools/call", + "params": { + "_meta": { + "progressToken": 5, + "vendor.example/custom": {"nested": ["a", 1, null]}, + "another-key": true + }, + "name": "my_tool" + } + }); + let req: CallToolRequest = serde_json::from_value(input.clone()).unwrap(); + let output = serde_json::to_value(&req).unwrap(); + assert_eq!(input, output); + } } diff --git a/crates/rmcp/src/model/task.rs b/crates/rmcp/src/model/task.rs index 8f49342..bbdd1c9 100644 --- a/crates/rmcp/src/model/task.rs +++ b/crates/rmcp/src/model/task.rs @@ -1,114 +1,90 @@ +//! Task types for the MCP Tasks extension (SEP-2663). +//! +//! Tasks are defined by the official `io.modelcontextprotocol/tasks` extension. +//! A server may respond to a supported request (currently `tools/call`) with a +//! [`CreateTaskResult`] (`resultType: "task"`) instead of the standard result. +//! The client then polls `tasks/get`, answers in-task server-to-client requests +//! via `tasks/update`, and may signal cancellation via `tasks/cancel`. + use serde::{Deserialize, Serialize}; -use serde_json::Value; -use super::Meta; +use super::{InputRequests, JsonObject, MetaObject, ResultType}; -/// Metadata for augmenting a request with task execution (spec `TaskMetadata`). -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct TaskMetadata { - #[serde(skip_serializing_if = "Option::is_none")] - pub ttl: Option, -} +/// Extension identifier for the MCP Tasks extension (SEP-2663). +pub const TASKS_EXTENSION_ID: &str = "io.modelcontextprotocol/tasks"; -impl TaskMetadata { - pub fn new() -> Self { - Self::default() - } - - pub fn with_ttl(mut self, ttl: u64) -> Self { - self.ttl = Some(ttl); - self - } -} - -/// Metadata for associating messages with a task (spec `RelatedTaskMetadata`). -/// -/// Carried in `_meta` under the key `"io.modelcontextprotocol/related-task"`. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct RelatedTaskMetadata { - pub task_id: String, -} - -impl RelatedTaskMetadata { - pub fn new(task_id: impl Into) -> Self { - Self { - task_id: task_id.into(), - } - } - - /// The well-known `_meta` key for related-task metadata. - pub const META_KEY: &str = "io.modelcontextprotocol/related-task"; -} - -/// Canonical task lifecycle status as defined by SEP-1686. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +/// Canonical task lifecycle status (SEP-2663). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] #[serde(rename_all = "snake_case")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub enum TaskStatus { - /// The receiver accepted the request and is currently working on it. + /// The request is currently being processed. #[default] Working, - /// The receiver requires additional input before work can continue. + /// The server needs input from the client before the task can proceed. InputRequired, - /// The underlying operation completed successfully and the result is ready. + /// The request completed successfully and the result is available. + /// This includes tool calls that returned results with `isError: true`. Completed, - /// The underlying operation failed and will not continue. + /// The request failed due to a JSON-RPC error during execution. Failed, - /// The task was cancelled and will not continue processing. + /// The request was cancelled before completion. Cancelled, } -/// Primary Task object that surfaces metadata during the task lifecycle. -/// -/// Per spec, `lastUpdatedAt` and `ttl` are required fields. -/// `ttl` is nullable (`null` means unlimited retention). +impl TaskStatus { + /// Returns `true` for terminal statuses (`completed`, `failed`, `cancelled`). + pub fn is_terminal(&self) -> bool { + matches!(self, Self::Completed | Self::Failed | Self::Cancelled) + } +} + +/// Operational metadata about ongoing work (spec `Task`, SEP-2663). #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub struct Task { - /// Unique task identifier generated by the receiver. + /// Stable identifier for this task, generated by the server. pub task_id: String, - /// Current lifecycle status (see [`TaskStatus`]). + /// Current task status. pub status: TaskStatus, - /// Optional human-readable status message for UI surfaces. + /// Optional message describing the current task state. + /// This MAY be exposed to the end-user or model. #[serde(skip_serializing_if = "Option::is_none")] pub status_message: Option, - /// ISO-8601 creation timestamp. + /// ISO 8601 timestamp when the task was created. pub created_at: String, - /// ISO-8601 timestamp for the most recent status change. + /// ISO 8601 timestamp when the task was last updated. pub last_updated_at: String, - /// Retention window in milliseconds that the receiver agreed to honor. - /// `None` (serialized as `null`) means unlimited retention. - pub ttl: Option, - /// Suggested polling interval (milliseconds). + /// Time-to-live duration from creation in integer milliseconds; `None` + /// (serialized as `null`) means unlimited. The server may discard the task + /// after the TTL elapses. This value MAY change over the lifetime of a task. + pub ttl_ms: Option, + /// Suggested polling interval in integer milliseconds. Clients SHOULD honor + /// this value to avoid overwhelming the server. This value MAY change over + /// the lifetime of a task. #[serde(skip_serializing_if = "Option::is_none")] - pub poll_interval: Option, + pub poll_interval_ms: Option, } impl Task { - /// Create a new Task with required fields. + /// Create a new task with required fields. pub fn new( - task_id: String, + task_id: impl Into, status: TaskStatus, - created_at: String, - last_updated_at: String, + created_at: impl Into, + last_updated_at: impl Into, ) -> Self { Self { - task_id, + task_id: task_id.into(), status, status_message: None, - created_at, - last_updated_at, - ttl: None, - poll_interval: None, + created_at: created_at.into(), + last_updated_at: last_updated_at.into(), + ttl_ms: None, + poll_interval_ms: None, } } @@ -119,116 +95,450 @@ impl Task { } /// Set the TTL in milliseconds. `None` means unlimited retention. - pub fn with_ttl(mut self, ttl: u64) -> Self { - self.ttl = Some(ttl); + pub fn with_ttl_ms(mut self, ttl_ms: u64) -> Self { + self.ttl_ms = Some(ttl_ms); self } - /// Set the poll interval in milliseconds. - pub fn with_poll_interval(mut self, poll_interval: u64) -> Self { - self.poll_interval = Some(poll_interval); + /// Set the suggested poll interval in milliseconds. + pub fn with_poll_interval_ms(mut self, poll_interval_ms: u64) -> Self { + self.poll_interval_ms = Some(poll_interval_ms); self } } -/// Wrapper returned by task-augmented requests (CreateTaskResult in SEP-1686). -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +/// Status-specific payload carried alongside the base [`Task`] fields in a +/// [`DetailedTask`]. +/// +/// Mirrors the spec's `WorkingTask` / `InputRequiredTask` / `CompletedTask` / +/// `FailedTask` / `CancelledTask` union: the variant is discriminated by the +/// `status` field on the wire, with the payload fields inlined at the top level. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub enum TaskPayload { + /// `status: "working"` — no additional payload. + Working, + /// `status: "input_required"` — outstanding server-to-client requests. + InputRequired { + /// Server-to-client requests that need to be fulfilled during task + /// execution. Keys are arbitrary identifiers for matching requests + /// to responses, unique over the lifetime of the task. + input_requests: InputRequests, + }, + /// `status: "completed"` — the final result of the task. The structure + /// matches the result type of the original request (e.g. `CallToolResult`). + Completed { + /// The final result of the original request. + result: JsonObject, + }, + /// `status: "failed"` — the JSON-RPC error that caused the task to fail. + Failed { + /// The JSON-RPC error object. + error: JsonObject, + }, + /// `status: "cancelled"` — no additional payload. + Cancelled, +} + +impl TaskPayload { + /// The [`TaskStatus`] this payload corresponds to. + pub fn status(&self) -> TaskStatus { + match self { + Self::Working => TaskStatus::Working, + Self::InputRequired { .. } => TaskStatus::InputRequired, + Self::Completed { .. } => TaskStatus::Completed, + Self::Failed { .. } => TaskStatus::Failed, + Self::Cancelled => TaskStatus::Cancelled, + } + } +} + +/// A task with its status-specific payload inlined (spec `DetailedTask`). +/// +/// Used by `tasks/get` responses ([`GetTaskResult`]) and `notifications/tasks` +/// ([`TaskStatusNotificationParams`](crate::model::TaskStatusNotificationParams)). +/// On the wire, the payload fields (`inputRequests` / `result` / `error`) are +/// flattened at the top level next to the base [`Task`] fields, and `status` +/// discriminates the variant. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub struct DetailedTask { + /// Base task metadata. Its `status` always agrees with the payload. + pub task: Task, + /// Status-specific payload. + pub payload: TaskPayload, +} + +impl DetailedTask { + /// Build a `DetailedTask`, forcing `task.status` to match the payload. + pub fn new(mut task: Task, payload: TaskPayload) -> Self { + task.status = payload.status(); + Self { task, payload } + } + + /// The current status. + pub fn status(&self) -> TaskStatus { + self.task.status + } +} + +// Wire shape helper: base Task fields + optional payload fields, all flattened. +#[derive(Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +struct DetailedTaskWire { + #[serde(flatten)] + task: Task, + #[serde(skip_serializing_if = "Option::is_none")] + input_requests: Option, + #[serde(skip_serializing_if = "Option::is_none")] + result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +impl From for DetailedTaskWire { + fn from(value: DetailedTask) -> Self { + let DetailedTask { task, payload } = value; + let (input_requests, result, error) = match payload { + TaskPayload::Working | TaskPayload::Cancelled => (None, None, None), + TaskPayload::InputRequired { input_requests } => (Some(input_requests), None, None), + TaskPayload::Completed { result } => (None, Some(result), None), + TaskPayload::Failed { error } => (None, None, Some(error)), + }; + Self { + task, + input_requests, + result, + error, + } + } +} + +impl TryFrom for DetailedTask { + type Error = String; + fn try_from(wire: DetailedTaskWire) -> Result { + let payload = match wire.task.status { + TaskStatus::Working => TaskPayload::Working, + TaskStatus::Cancelled => TaskPayload::Cancelled, + TaskStatus::InputRequired => TaskPayload::InputRequired { + input_requests: wire.input_requests.ok_or_else(|| { + "task with status \"input_required\" is missing `inputRequests`".to_owned() + })?, + }, + TaskStatus::Completed => TaskPayload::Completed { + result: wire.result.ok_or_else(|| { + "task with status \"completed\" is missing `result`".to_owned() + })?, + }, + TaskStatus::Failed => TaskPayload::Failed { + error: wire + .error + .ok_or_else(|| "task with status \"failed\" is missing `error`".to_owned())?, + }, + }; + Ok(DetailedTask { + task: wire.task, + payload, + }) + } +} + +impl Serialize for DetailedTask { + fn serialize(&self, serializer: S) -> Result { + DetailedTaskWire::from(self.clone()).serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for DetailedTask { + fn deserialize>(deserializer: D) -> Result { + let wire = DetailedTaskWire::deserialize(deserializer)?; + Self::try_from(wire).map_err(serde::de::Error::custom) + } +} + +#[cfg(feature = "schemars")] +impl schemars::JsonSchema for DetailedTask { + fn schema_name() -> std::borrow::Cow<'static, str> { + "DetailedTask".into() + } + fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema { + // The actual wire shape: base Task fields plus the optional + // status-specific payload fields (inputRequests / result / error). + ::json_schema(generator) + } +} + +/// Result returned in lieu of a standard result to indicate the request will +/// be processed asynchronously (spec `CreateTaskResult`, `resultType: "task"`). +/// +/// The embedded task is the seed state for the task; the client uses +/// `task.task_id` for all subsequent `tasks/get`, `tasks/update`, and +/// `tasks/cancel` calls. +#[derive(Debug, Clone, PartialEq, Serialize)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub struct CreateTaskResult { + /// Always `"task"`. + pub result_type: ResultType, + /// Seed state of the newly created task, flattened at the top level. + #[serde(flatten)] pub task: Task, - #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, + #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] + pub meta: Option, +} + +// Custom deserializer that requires `resultType: "task"`. Without this, +// `CreateTaskResult` would greedily match other task-shaped results (e.g. +// `tasks/get` responses, which also carry `taskId`/`status` at the top level +// but use `resultType: "complete"`) inside `#[serde(untagged)]` unions such +// as `ServerResult`. +impl<'de> Deserialize<'de> for CreateTaskResult { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + #[derive(Deserialize)] + #[serde(rename_all = "camelCase")] + struct Helper { + result_type: ResultType, + #[serde(flatten)] + task: Task, + #[serde(rename = "_meta", default)] + meta: Option, + } + let helper = Helper::deserialize(deserializer)?; + if !helper.result_type.is_task() { + return Err(serde::de::Error::custom( + "CreateTaskResult requires resultType to be \"task\"", + )); + } + Ok(CreateTaskResult { + result_type: helper.result_type, + task: helper.task, + meta: helper.meta, + }) + } } impl CreateTaskResult { - /// Create a new CreateTaskResult. + /// Create a new `CreateTaskResult` from the seed task state. pub fn new(task: Task) -> Self { - Self { task, meta: None } + Self { + result_type: ResultType::TASK, + task, + meta: None, + } } /// Sets the protocol-level metadata for this result. - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } } -/// Response to a `tasks/get` request. +/// Response to a `tasks/get` request (spec `GetTaskResult = Result & DetailedTask`). /// -/// Per spec, `GetTaskResult = allOf[Result, Task]` — the Task fields are -/// flattened at the top level, not nested under a `task` key. +/// `resultType` is `"complete"` — this is the standard result shape for +/// `tasks/get`, not a task handle. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub struct GetTaskResult { + /// Result type discriminator. `tasks/get` responses are standard results: + /// `"complete"` (SEP-2322). Absent values deserialize as `"complete"`. + #[serde(default)] + pub result_type: ResultType, #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, + pub meta: Option, + /// The task with status-specific payload inlined. #[serde(flatten)] - pub task: Task, + pub task: DetailedTask, } impl GetTaskResult { - pub fn new(task: Task) -> Self { - Self { meta: None, task } + pub fn new(task: DetailedTask) -> Self { + Self { + result_type: ResultType::COMPLETE, + meta: None, + task, + } } } -/// Response to a `tasks/result` request. +/// Empty acknowledgement for `tasks/update` and `tasks/cancel` (SEP-2663). /// -/// Per spec, the result structure matches the original request type -/// (e.g., `CallToolResult` for `tools/call`). This is represented as -/// an open object. The payload is the original request's result -/// serialized as a JSON value. +/// The spec requires these acks to be empty results carrying the SEP-2322 +/// `resultType: "complete"` discriminator; task state changes are observed +/// via the next `tasks/get`. #[derive(Debug, Clone, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] -pub struct GetTaskPayloadResult(pub Value); - -impl GetTaskPayloadResult { - /// Create a new GetTaskPayloadResult with the given value. - pub fn new(value: Value) -> Self { - Self(value) - } +pub struct TaskAckResult { + /// Always `"complete"`. + pub result_type: ResultType, + #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] + pub meta: Option, } -// Custom Deserialize that always fails, so that `GetTaskPayloadResult` is skipped -// during `#[serde(untagged)]` enum deserialization (e.g. `ServerResult`). -// The payload has the same JSON shape as `CustomResult(Value)`, so they are -// indistinguishable. `CustomResult` acts as the catch-all instead. -// `GetTaskPayloadResult` should be constructed programmatically via `::new()`. -impl<'de> serde::Deserialize<'de> for GetTaskPayloadResult { +// Custom deserializer that requires `resultType: "complete"` and rejects any +// other fields. Without this, `TaskAckResult` would greedily match arbitrary +// result objects carrying a `resultType` key inside `#[serde(untagged)]` +// unions such as `ServerResult`, shadowing `CustomResult` and losing data. +impl<'de> Deserialize<'de> for TaskAckResult { fn deserialize(deserializer: D) -> Result where D: serde::Deserializer<'de>, { - // Consume the value so the deserializer state stays consistent. - serde::de::IgnoredAny::deserialize(deserializer)?; - Err(serde::de::Error::custom( - "GetTaskPayloadResult cannot be deserialized directly; \ - use CustomResult as the catch-all", - )) + #[derive(Deserialize)] + #[serde(rename_all = "camelCase", deny_unknown_fields)] + struct Helper { + result_type: ResultType, + #[serde(rename = "_meta", default)] + meta: Option, + } + let helper = Helper::deserialize(deserializer)?; + if !helper.result_type.is_complete() { + return Err(serde::de::Error::custom( + "TaskAckResult requires resultType to be \"complete\"", + )); + } + Ok(TaskAckResult { + result_type: helper.result_type, + meta: helper.meta, + }) } } -/// Response to a `tasks/cancel` request. -/// -/// Per spec, `CancelTaskResult = allOf[Result, Task]` — same shape as `GetTaskResult`. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct CancelTaskResult { - #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] - pub meta: Option, - #[serde(flatten)] - pub task: Task, +impl Default for TaskAckResult { + fn default() -> Self { + Self { + result_type: ResultType::COMPLETE, + meta: None, + } + } } -impl CancelTaskResult { - pub fn new(task: Task) -> Self { - Self { meta: None, task } +impl TaskAckResult { + pub fn new() -> Self { + Self::default() + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn base_task(status: TaskStatus) -> Task { + Task::new( + "task-1", + status, + "2025-11-25T10:30:00Z", + "2025-11-25T10:40:00Z", + ) + .with_ttl_ms(60000) + .with_poll_interval_ms(5000) + } + + #[test] + fn create_task_result_wire_shape() { + let result = CreateTaskResult::new(base_task(TaskStatus::Working)); + let value = serde_json::to_value(&result).unwrap(); + assert_eq!( + value, + json!({ + "resultType": "task", + "taskId": "task-1", + "status": "working", + "createdAt": "2025-11-25T10:30:00Z", + "lastUpdatedAt": "2025-11-25T10:40:00Z", + "ttlMs": 60000, + "pollIntervalMs": 5000 + }) + ); + let roundtrip: CreateTaskResult = serde_json::from_value(value).unwrap(); + assert_eq!(roundtrip, result); + } + + #[test] + fn ttl_ms_null_means_unlimited() { + let mut task = base_task(TaskStatus::Working); + task.ttl_ms = None; + let value = serde_json::to_value(&task).unwrap(); + assert_eq!(value["ttlMs"], serde_json::Value::Null); + let roundtrip: Task = serde_json::from_value(value).unwrap(); + assert_eq!(roundtrip.ttl_ms, None); + } + + #[test] + fn detailed_task_completed_roundtrip() { + let detailed = DetailedTask::new( + base_task(TaskStatus::Working), + TaskPayload::Completed { + result: serde_json::from_value(json!({ + "content": [{"type": "text", "text": "ok"}], + "isError": false + })) + .unwrap(), + }, + ); + // Status is forced to match the payload. + assert_eq!(detailed.status(), TaskStatus::Completed); + let value = serde_json::to_value(&detailed).unwrap(); + assert_eq!(value["status"], "completed"); + assert_eq!(value["result"]["isError"], false); + let roundtrip: DetailedTask = serde_json::from_value(value).unwrap(); + assert_eq!(roundtrip, detailed); + } + + #[test] + fn detailed_task_input_required_requires_input_requests() { + let err = serde_json::from_value::(json!({ + "taskId": "task-1", + "status": "input_required", + "createdAt": "2025-11-25T10:30:00Z", + "lastUpdatedAt": "2025-11-25T10:40:00Z", + "ttlMs": null + })) + .unwrap_err(); + assert!(err.to_string().contains("inputRequests")); + } + + #[test] + fn detailed_task_failed_roundtrip() { + let detailed = DetailedTask::new( + base_task(TaskStatus::Failed), + TaskPayload::Failed { + error: serde_json::from_value(json!({ + "code": -32603, + "message": "boom" + })) + .unwrap(), + }, + ); + let value = serde_json::to_value(&detailed).unwrap(); + assert_eq!(value["status"], "failed"); + assert_eq!(value["error"]["code"], -32603); + let roundtrip: DetailedTask = serde_json::from_value(value).unwrap(); + assert_eq!(roundtrip, detailed); + } + + #[test] + fn get_task_result_flattens_detailed_task() { + let result = GetTaskResult::new(DetailedTask::new( + base_task(TaskStatus::Working), + TaskPayload::Working, + )); + let value = serde_json::to_value(&result).unwrap(); + assert_eq!(value["taskId"], "task-1"); + assert_eq!(value["status"], "working"); + let roundtrip: GetTaskResult = serde_json::from_value(value).unwrap(); + assert_eq!(roundtrip, result); } } diff --git a/crates/rmcp/src/model/tool.rs b/crates/rmcp/src/model/tool.rs index 2ed89d6..be7cbbd 100644 --- a/crates/rmcp/src/model/tool.rs +++ b/crates/rmcp/src/model/tool.rs @@ -7,7 +7,7 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use serde_json::Value; -use super::{Icon, JsonObject, Meta}; +use super::{Icon, JsonObject, MetaObject}; /// A tool that can be used by a model. #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -31,71 +31,12 @@ pub struct Tool { #[serde(skip_serializing_if = "Option::is_none")] /// Optional additional tool information. pub annotations: Option, - /// Execution-related configuration including task support mode. - #[serde(skip_serializing_if = "Option::is_none")] - pub execution: Option, /// Optional list of icons for the tool #[serde(skip_serializing_if = "Option::is_none")] pub icons: Option>, /// Optional additional metadata for this tool #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - pub meta: Option, -} - -/// Per-tool task support mode as defined in the MCP specification. -/// -/// This enum indicates whether a tool supports task-based invocation, -/// allowing clients to know how to properly call the tool. -/// -/// See [Tool-Level Negotiation](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks#tool-level-negotiation). -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] -#[serde(rename_all = "lowercase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[expect(clippy::exhaustive_enums, reason = "intentionally exhaustive")] -pub enum TaskSupport { - /// Clients MUST NOT invoke this tool as a task (default behavior). - #[default] - Forbidden, - /// Clients MAY invoke this tool as either a task or a normal call. - Optional, - /// Clients MUST invoke this tool as a task. - Required, -} - -/// Execution-related configuration for a tool. -/// -/// This struct contains settings that control how a tool should be executed, -/// including task support configuration. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] -#[serde(rename_all = "camelCase")] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[non_exhaustive] -pub struct ToolExecution { - /// Indicates whether this tool supports task-based invocation. - /// - /// When not present or set to `Forbidden`, clients MUST NOT invoke this tool as a task. - /// When set to `Optional`, clients MAY invoke this tool as a task or normal call. - /// When set to `Required`, clients MUST invoke this tool as a task. - #[serde(skip_serializing_if = "Option::is_none")] - pub task_support: Option, -} - -impl ToolExecution { - /// Create a new empty ToolExecution configuration. - pub fn new() -> Self { - Self::default() - } - - /// Create a ToolExecution from raw optional fields. - pub fn from_raw(task_support: Option) -> Self { - Self { task_support } - } - - /// Set the task support mode. - pub fn with_task_support(mut self, task_support: TaskSupport) -> Self { - self.task_support = Some(task_support); - self - } + pub meta: Option, } /// Additional properties describing a Tool to clients. @@ -232,7 +173,6 @@ impl Tool { input_schema: input_schema.into(), output_schema: None, annotations: None, - execution: None, icons: None, meta: None, } @@ -255,7 +195,6 @@ impl Tool { input_schema: input_schema.into(), output_schema: None, annotations: None, - execution: None, icons: None, meta: None, } @@ -286,7 +225,7 @@ impl Tool { } /// Set the metadata - pub fn with_meta(mut self, meta: Meta) -> Self { + pub fn with_meta(mut self, meta: MetaObject) -> Self { self.meta = Some(meta); self } @@ -298,32 +237,10 @@ impl Tool { } } - /// Set the execution configuration for this tool. - pub fn with_execution(mut self, execution: ToolExecution) -> Self { - self.execution = Some(execution); - self - } - - /// Returns the task support mode for this tool. - /// - /// Returns `TaskSupport::Forbidden` if not explicitly set. - pub fn task_support(&self) -> TaskSupport { - self.execution - .as_ref() - .and_then(|e| e.task_support) - .unwrap_or_default() - } - /// Set the output schema using a type that implements JsonSchema - /// - /// # Panics - /// - /// Panics if the generated schema does not have root type "object" as required by MCP specification. #[cfg(feature = "server")] pub fn with_output_schema(mut self) -> Self { - let schema = crate::handler::server::tool::schema_for_output::() - .unwrap_or_else(|e| panic!("Invalid output schema for tool '{}': {}", self.name, e)); - self.output_schema = Some(schema); + self.output_schema = Some(crate::handler::server::tool::schema_for_output::()); self } diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs index 29d822a..63b38e4 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -1,3 +1,5 @@ +use std::{borrow::Cow, sync::OnceLock}; + use futures::FutureExt; #[cfg(not(feature = "local"))] use futures::future::BoxFuture; @@ -51,9 +53,10 @@ use crate::model::ServerNotification; use crate::{ error::ErrorData as McpError, model::{ - CancelledNotification, CancelledNotificationParam, Extensions, GetExtensions, GetMeta, - JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Meta, - NumberOrString, ProgressToken, RequestId, + CancelledNotification, CancelledNotificationParam, ClientCapabilities, Extensions, + GetExtensions, GetMeta, Implementation, JsonRpcError, JsonRpcMessage, JsonRpcNotification, + JsonRpcRequest, JsonRpcResponse, NotificationMetaObject, NumberOrString, ProgressToken, + ProtocolVersion, RequestId, RequestMetaObject, }, transport::{DynamicTransportError, IntoTransport, Transport}, }; @@ -82,10 +85,15 @@ pub enum ServiceError { TransportClosed, #[error("Unexpected response type")] UnexpectedResponse, + #[error("subscription consumer lagged behind its {capacity}-message buffer")] + SubscriptionLagged { capacity: usize }, #[error("task cancelled for reason {}", reason.as_deref().unwrap_or(""))] Cancelled { reason: Option }, #[error("request timeout after {}", chrono::Duration::from_std(*timeout).unwrap_or_default())] Timeout { timeout: Duration }, + /// The peer kept returning `input_required` beyond the configured round cap. + #[error("input_required did not complete within {max_rounds} MRTR rounds")] + InputRequiredRoundsExceeded { max_rounds: usize }, } trait TransferObject: @@ -106,22 +114,166 @@ impl TransferObject for T where #[allow(private_bounds, reason = "there's no the third implementation")] pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone { - type Req: TransferObject + GetMeta + GetExtensions; + type Req: TransferObject + GetMeta + GetExtensions; type Resp: TransferObject; type Not: TryInto + From + TransferObject; - type PeerReq: TransferObject + GetMeta + GetExtensions; + type PeerReq: TransferObject + GetMeta + GetExtensions; type PeerResp: TransferObject; type PeerNot: TryInto + From + TransferObject - + GetMeta + + GetMeta + GetExtensions; type InitializeError; const IS_CLIENT: bool; type Info: TransferObject; type PeerInfo: TransferObject; + #[doc(hidden)] + fn configure_direct_peer(_peer: &Peer, _info: &Self::Info) {} + #[doc(hidden)] + fn peer_cancelled_params(_notification: &Self::PeerNot) -> Option<&CancelledNotificationParam> { + None + } + #[doc(hidden)] + fn peer_cancels_subscriptions(_peer: &Peer) -> bool { + false + } + #[doc(hidden)] + fn is_subscription_request(_request: &Self::Req) -> bool { + false + } + /// Invalidate any response cache affected by an inbound peer notification. + /// + /// The serve loop calls this for every notification *before* subscription + /// routing, so cache invalidation still runs when a notification is + /// delivered through a `listen` subscription channel rather than the + /// [`Service::handle_notification`] callbacks. + #[doc(hidden)] + fn invalidate_response_cache( + _peer: &Peer, + _notification: &Self::PeerNot, + ) -> impl Future + MaybeSendFuture { + async {} + } + + #[doc(hidden)] + fn enforce_request_association( + _request: &Self::Req, + _peer_info: Option<&Self::PeerInfo>, + _in_request_handler_scope: bool, + ) -> Result<(), ServiceError> { + Ok(()) + } + + /// Receive-side counterpart of [`Self::enforce_request_association`]: + /// SEP-2260 says clients receiving a server-to-client request with no + /// associated outbound request should reject it with invalid params. An + /// error return is sent back to the peer instead of dispatching to the + /// handler. + #[doc(hidden)] + fn enforce_peer_request_association( + _peer_request: &Self::PeerReq, + _peer_info: Option<&Self::PeerInfo>, + _association: PeerRequestAssociation, + ) -> Result<(), McpError> { + Ok(()) + } +} + +/// How an inbound peer request relates to this side's in-flight outbound +/// requests (SEP-2260). +/// +/// SEP-2260 defines no wire field for association, so only stream-separating +/// transports (streamable HTTP) can observe it; other transports yield +/// [`Self::Unknown`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[expect(clippy::exhaustive_enums, reason = "intentionally exhaustive")] +pub enum PeerRequestAssociation { + /// Arrived on the response stream of an in-flight outbound request. + Associated, + /// Arrived on a stream tied to no in-flight outbound request (e.g. the + /// streamable HTTP standalone GET stream). + Unassociated, + /// The transport cannot distinguish streams; only the coarse in-flight + /// signal is available. + Unknown { has_pending_outbound_request: bool }, +} + +/// Whether `version` predates `2026-07-28`, the revision that replaced the +/// `initialize` handshake with per-request metadata. +pub(crate) fn is_legacy_version(version: &ProtocolVersion) -> bool { + version.as_str() < ProtocolVersion::V_2026_07_28.as_str() +} + +pub(crate) fn uses_legacy_lifecycle( + protocol_version: Option<&ProtocolVersion>, + uses_discover_lifecycle: bool, +) -> bool { + !uses_discover_lifecycle && protocol_version.is_none_or(is_legacy_version) +} + +pub(crate) fn peer_request_association( + request: &Req, + local_responder_pool: &std::collections::HashMap, +) -> PeerRequestAssociation { + match request.extensions().get::() { + None => PeerRequestAssociation::Unknown { + has_pending_outbound_request: !local_responder_pool.is_empty(), + }, + Some(InboundStreamOrigin::Unassociated) => PeerRequestAssociation::Unassociated, + Some(InboundStreamOrigin::OutboundRequest(id)) => { + if local_responder_pool.contains_key(id) { + PeerRequestAssociation::Associated + } else { + PeerRequestAssociation::Unassociated + } + } + } +} + +tokio::task_local! { + pub(crate) static ORIGINATING_REQUEST: RequestId; +} + +pub(crate) fn in_request_handler_scope() -> bool { + ORIGINATING_REQUEST.try_with(|_| ()).is_ok() +} + +/// Marker in an outbound request's non-serialized [`Extensions`] identifying +/// the in-flight peer request it was issued from (SEP-2260). Attached for both +/// roles whenever a request is sent from within a request handler; the +/// streamable HTTP server reads it to deliver server-initiated requests on the +/// originating request's SSE stream. Never on the wire (SEP-2260 defines no +/// wire field), so session managers that serialize messages between processes +/// lose it and such requests fall back to the standalone stream with a warning. +/// +/// # Caller requirements +/// +/// From protocol version `2026-07-28`, server-to-client sampling, roots, and +/// elicitation requests must be issued while handling a client request; +/// outside a handler they return an `invalid_request` error. The association +/// is task-local and does not cross `tokio::spawn`, so use the task manager +/// for long-running work. +/// +/// The client receive-side mirror is [`InboundStreamOrigin`]. +#[derive(Debug, Clone, PartialEq, Eq)] +#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] +pub struct OriginatingRequestId(pub RequestId); + +/// Marker in an inbound request's non-serialized [`Extensions`] recording +/// which HTTP response stream it arrived on: the receive-side mirror of +/// [`OriginatingRequestId`]. Never on the wire (SEP-2260 defines no wire +/// field); when absent, the coarse in-flight check applies. +#[derive(Debug, Clone, PartialEq, Eq)] +#[expect(clippy::exhaustive_enums, reason = "intentionally exhaustive")] +pub enum InboundStreamOrigin { + /// The standalone GET stream, or a POST response stream not tied to an + /// outbound request. + Unassociated, + /// The SSE response stream of the POST that carried this outbound request. + OutboundRequest(RequestId), } pub type TxJsonRpcMessage = @@ -145,6 +297,19 @@ pub trait Service: Send + Sync + 'static { context: NotificationContext, ) -> impl Future> + MaybeSendFuture + '_; fn get_info(&self) -> R::Info; + /// The protocol versions this service can speak, bounding what `initialize` + /// negotiation may agree to. + /// + /// Servers normally override + /// [`ServerHandler::supported_protocol_versions`] instead of this method; + /// the blanket `Service` impl forwards to it. This method exists so the + /// transport and handshake layers, which see only a `Service`, can read the + /// list and avoid agreeing to a version the server cannot serve. + /// + /// [`ServerHandler::supported_protocol_versions`]: crate::handler::server::ServerHandler::supported_protocol_versions + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(ProtocolVersion::KNOWN_VERSIONS) + } } #[cfg(feature = "local")] @@ -160,6 +325,12 @@ pub trait Service: 'static { context: NotificationContext, ) -> impl Future> + MaybeSendFuture + '_; fn get_info(&self) -> R::Info; + /// The protocol versions this service can speak. + /// + /// See the non-`local` variant of this trait for details. + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(ProtocolVersion::KNOWN_VERSIONS) + } } pub trait ServiceExt: Service + Sized { @@ -211,6 +382,10 @@ impl Service for Box> { fn get_info(&self) -> R::Info { DynService::get_info(self.as_ref()) } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + DynService::supported_protocol_versions(self.as_ref()) + } } #[cfg(not(feature = "local"))] @@ -226,6 +401,10 @@ pub trait DynService: Send + Sync { context: NotificationContext, ) -> MaybeBoxFuture<'_, Result<(), McpError>>; fn get_info(&self) -> R::Info; + /// See [`Service::supported_protocol_versions`]. + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(ProtocolVersion::KNOWN_VERSIONS) + } } #[cfg(feature = "local")] @@ -241,6 +420,10 @@ pub trait DynService { context: NotificationContext, ) -> MaybeBoxFuture<'_, Result<(), McpError>>; fn get_info(&self) -> R::Info; + /// See [`Service::supported_protocol_versions`]. + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(ProtocolVersion::KNOWN_VERSIONS) + } } impl> DynService for S { @@ -261,6 +444,9 @@ impl> DynService for S { fn get_info(&self) -> R::Info { self.get_info() } + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Service::supported_protocol_versions(self) + } } use std::{ @@ -283,6 +469,17 @@ pub trait ProgressTokenProvider: Send + Sync + 'static { pub type AtomicU32RequestIdProvider = AtomicU32Provider; pub type AtomicU32ProgressTokenProvider = AtomicU32Provider; +pub(crate) fn remove_pending_request( + pending_requests: &mut HashMap, + response_id: &RequestId, +) -> Option { + pending_requests.remove(response_id).or_else(|| { + response_id + .numeric_string_value() + .and_then(|id| pending_requests.remove(&RequestId::Number(id))) + }) +} + #[derive(Debug, Default)] pub struct AtomicU32Provider { id: AtomicU64, @@ -333,7 +530,13 @@ impl ProgressNotificationToken for ServerNotification { } type Responder = tokio::sync::oneshot::Sender; +struct PendingResponse { + responder: Responder, + is_subscription: bool, +} type ProgressTimeoutWatchers = Arc>>>; +type SubscriptionChannel = (mpsc::Sender, usize); +type SubscriptionChannelMap = HashMap>; /// A handle to a remote request /// @@ -394,10 +597,12 @@ impl RequestHandle { has_progress_reset_rx, ) .await; + self.peer.unregister_subscription(&self.id); result } async fn send_timeout_cancel_notification(&self, reason: &str) { + self.peer.unregister_subscription(&self.id); let notification = CancelledNotification { params: CancelledNotificationParam { request_id: Some(self.id.clone()), @@ -454,11 +659,10 @@ impl RequestHandle { None => None, } }, if reset_timeout_on_progress && idle_sleep.is_some() && self.progress_reset_rx.is_some() => { - if progress.is_some() { - if let Some((timeout, sleep)) = idle_sleep.as_mut() { + if progress.is_some() + && let Some((timeout, sleep)) = idle_sleep.as_mut() { sleep.as_mut().reset(tokio::time::Instant::now() + *timeout); } - } } } } @@ -472,6 +676,7 @@ impl RequestHandle { self.progress_reset_rx.is_some(), ) .await; + self.peer.unregister_subscription(&self.id); let notification = CancelledNotification { params: CancelledNotificationParam { request_id: Some(self.id), @@ -512,18 +717,49 @@ pub(crate) enum PeerSinkMessage { }, } +#[derive(Debug, Clone)] +pub(crate) struct ClientRequestMetadata { + pub protocol_version: ProtocolVersion, + pub client_info: Implementation, + pub client_capabilities: ClientCapabilities, +} + /// An interface to fetch the remote client or server /// /// For general purpose, call [`Peer::send_request`] or [`Peer::send_notification`] to send message to remote peer. /// /// To create a cancellable request, call [`Peer::send_request_with_option`]. -#[derive(Clone)] pub struct Peer { tx: mpsc::Sender>, request_id_provider: Arc, progress_token_provider: Arc, progress_timeout_watchers: ProgressTimeoutWatchers, info: Arc>>>, + client_request_metadata: Arc>, + request_metadata_required: Arc, + subscription_channels: Arc>>, + #[cfg(feature = "client")] + response_cache: client::cache::PeerResponseCache, +} + +impl Clone for Peer +where + R::PeerInfo: Clone, +{ + fn clone(&self) -> Peer { + Self { + tx: self.tx.clone(), + request_id_provider: self.request_id_provider.clone(), + progress_token_provider: self.progress_token_provider.clone(), + progress_timeout_watchers: self.progress_timeout_watchers.clone(), + info: self.info.clone(), + client_request_metadata: self.client_request_metadata.clone(), + request_metadata_required: self.request_metadata_required.clone(), + subscription_channels: self.subscription_channels.clone(), + #[cfg(feature = "client")] + response_cache: self.response_cache.clone(), + } + } } impl std::fmt::Debug for Peer { @@ -541,7 +777,7 @@ type ProxyOutbound = mpsc::Receiver>; #[non_exhaustive] pub struct PeerRequestOptions { pub timeout: Option, - pub meta: Option, + pub meta: Option, /// Reset the request timeout when a matching progress notification is received. pub reset_timeout_on_progress: bool, /// Maximum total time to wait for the request, regardless of progress notifications. @@ -560,6 +796,14 @@ impl PeerRequestOptions { } } + /// Adds request metadata while preserving any other configured options. + /// + /// Explicit values take precedence over discover-lifecycle metadata defaults. + pub fn with_meta(mut self, meta: RequestMetaObject) -> Self { + self.meta = Some(meta); + self + } + pub fn reset_timeout_on_progress(mut self) -> Self { self.reset_timeout_on_progress = true; self @@ -585,6 +829,11 @@ impl Peer { progress_token_provider: Arc::new(AtomicU32ProgressTokenProvider::default()), progress_timeout_watchers: Default::default(), info: Arc::new(std::sync::RwLock::new(peer_info.map(Arc::new))), + client_request_metadata: Default::default(), + request_metadata_required: Default::default(), + subscription_channels: Default::default(), + #[cfg(feature = "client")] + response_cache: Default::default(), }, rx, ) @@ -617,11 +866,37 @@ impl Peer { pub async fn send_request_with_option( &self, - mut request: R::Req, + request: R::Req, options: PeerRequestOptions, ) -> Result, ServiceError> { + self.send_request_with_option_and_subscription(request, options, None) + .await + } + + async fn send_request_with_option_and_subscription( + &self, + mut request: R::Req, + options: PeerRequestOptions, + subscription_sender: Option>, + ) -> Result, ServiceError> { + R::enforce_request_association( + &request, + self.peer_info().as_deref(), + in_request_handler_scope(), + )?; + if let Ok(originating) = ORIGINATING_REQUEST.try_with(|id| id.clone()) { + request + .extensions_mut() + .insert(OriginatingRequestId(originating)); + } let id = self.request_id_provider.next_request_id(); let progress_token = self.progress_token_provider.next_progress_token(); + if let Some(metadata) = self.client_request_metadata.get() { + let meta = request.get_meta_mut(); + meta.set_protocol_version(metadata.protocol_version.clone()); + meta.set_client_info(metadata.client_info.clone()); + meta.set_client_capabilities(metadata.client_capabilities.clone()); + } if let Some(meta) = options.meta.clone() { request.get_meta_mut().extend(meta); } @@ -639,6 +914,10 @@ impl Peer { } else { None }; + if let Some(channel) = subscription_sender { + self.subscription_channels_write() + .insert(id.clone(), channel); + } if self .tx .send(PeerSinkMessage::Request { @@ -655,6 +934,7 @@ impl Peer { .await .remove(&progress_token); } + self.unregister_subscription(&id); return Err(ServiceError::TransportClosed); } Ok(RequestHandle { @@ -667,6 +947,67 @@ impl Peer { }) } + #[cfg(feature = "client")] + pub(crate) async fn send_subscription_request( + &self, + request: R::Req, + options: PeerRequestOptions, + channel_capacity: usize, + ) -> Result<(RequestHandle, mpsc::Receiver), ServiceError> { + let (sender, receiver) = mpsc::channel(channel_capacity); + let handle = self + .send_request_with_option_and_subscription( + request, + options, + Some((sender, channel_capacity)), + ) + .await?; + Ok((handle, receiver)) + } + + fn subscription_sender(&self, id: &RequestId) -> Option> { + self.subscription_channels_read().get(id).cloned() + } + + pub(crate) fn unregister_subscription(&self, id: &RequestId) { + self.subscription_channels_write().remove(id); + } + + fn subscription_channels_read( + &self, + ) -> std::sync::RwLockReadGuard<'_, SubscriptionChannelMap> { + match self.subscription_channels.read() { + Ok(channels) => channels, + Err(poisoned) => poisoned.into_inner(), + } + } + + fn subscription_channels_write( + &self, + ) -> std::sync::RwLockWriteGuard<'_, SubscriptionChannelMap> { + match self.subscription_channels.write() { + Ok(channels) => channels, + Err(poisoned) => poisoned.into_inner(), + } + } + + pub(crate) fn try_cancel_request(&self, id: RequestId, reason: Option) { + let notification = CancelledNotification { + params: CancelledNotificationParam { + request_id: Some(id), + reason, + meta: None, + }, + method: crate::model::CancelledNotificationMethod, + extensions: Default::default(), + }; + let (responder, _receiver) = tokio::sync::oneshot::channel(); + let _ = self.tx.try_send(PeerSinkMessage::Notification { + notification: notification.into(), + responder, + }); + } + async fn notify_progress_timeout_watcher(&self, progress_token: &ProgressToken) { let sender = self .progress_timeout_watchers @@ -700,6 +1041,21 @@ impl Peer { *self.info.write().expect("peer info lock poisoned") = Some(Arc::new(info)); } + pub(crate) fn set_client_request_metadata(&self, metadata: ClientRequestMetadata) { + let result = self.client_request_metadata.set(metadata); + debug_assert!(result.is_ok(), "client request metadata set more than once"); + } + + pub(crate) fn require_request_metadata(&self) { + self.request_metadata_required + .store(true, std::sync::atomic::Ordering::Release); + } + + pub(crate) fn request_metadata_required(&self) -> bool { + self.request_metadata_required + .load(std::sync::atomic::Ordering::Acquire) + } + pub fn is_transport_closed(&self) -> bool { self.tx.is_closed() } @@ -859,7 +1215,7 @@ pub struct RequestContext { /// this token will be cancelled when the [`CancelledNotification`] is received. pub ct: CancellationToken, pub id: RequestId, - pub meta: Meta, + pub meta: RequestMetaObject, pub extensions: Extensions, /// An interface to fetch the remote client or server pub peer: Peer, @@ -871,7 +1227,7 @@ impl RequestContext { Self { ct: CancellationToken::new(), id, - meta: Meta::default(), + meta: RequestMetaObject::default(), extensions: Extensions::default(), peer, } @@ -880,11 +1236,35 @@ impl RequestContext { #[cfg(feature = "server")] impl RequestContext { - /// The protocol version the client negotiated, or `None` before peer info is recorded. + /// The current request's protocol version, falling back to legacy handshake state. pub fn protocol_version(&self) -> Option { - self.peer - .peer_info() - .map(|info| info.protocol_version.clone()) + self.meta.protocol_version().or_else(|| { + self.peer + .peer_info() + .map(|info| info.protocol_version.clone()) + }) + } + + /// The current request's client implementation, falling back only for legacy sessions. + pub fn client_info(&self) -> Option { + if self.peer.request_metadata_required() { + self.meta.client_info() + } else { + self.meta + .client_info() + .or_else(|| self.peer.peer_info().map(|info| info.client_info.clone())) + } + } + + /// The current request's client capabilities, falling back only for legacy sessions. + pub fn client_capabilities(&self) -> Option { + if self.peer.request_metadata_required() { + self.meta.client_capabilities() + } else { + self.meta + .client_capabilities() + .or_else(|| self.peer.peer_info().map(|info| info.capabilities.clone())) + } } } @@ -892,7 +1272,7 @@ impl RequestContext { #[derive(Debug, Clone)] #[non_exhaustive] pub struct NotificationContext { - pub meta: Meta, + pub meta: NotificationMetaObject, pub extensions: Extensions, /// An interface to fetch the remote client or server pub peer: Peer, @@ -929,7 +1309,15 @@ where E: std::error::Error + Send + Sync + 'static, { let (peer, peer_rx) = Peer::new(Arc::new(AtomicU32RequestIdProvider::default()), peer_info); - serve_inner(service, transport.into_transport(), peer, peer_rx, ct) + R::configure_direct_peer(&peer, &service.get_info()); + serve_inner( + service, + transport.into_transport(), + peer, + peer_rx, + VecDeque::new(), + ct, + ) } /// Spawn a task that may hold `!Send` state when the `local` feature is active. @@ -954,12 +1342,19 @@ where tokio::task::spawn_local(future) } +/// Run the service loop over `transport`. +/// +/// `initial_messages` are messages the caller already read from `transport` +/// (e.g. the first request of an `initialize`-less session). They are +/// dispatched by the loop, in order, before anything else is read from the +/// transport, so that their handlers run with the loop draining `peer_rx`. #[instrument(skip_all)] fn serve_inner( service: S, transport: T, peer: Peer, mut peer_rx: tokio::sync::mpsc::Receiver>, + initial_messages: VecDeque>, ct: CancellationToken, ) -> RunningService where @@ -979,7 +1374,7 @@ where } let mut local_responder_pool = - HashMap::>>::new(); + HashMap::>>::new(); let mut local_ct_pool = HashMap::::new(); let shared_service = Arc::new(service); // for return @@ -992,7 +1387,7 @@ where let current_span = tracing::Span::current(); let handle = spawn_service_task(async move { let mut transport = transport.into_transport(); - let mut batch_messages = VecDeque::>::new(); + let mut batch_messages = initial_messages; let mut send_task_set = tokio::task::JoinSet::::new(); let mut response_send_tasks = tokio::task::JoinSet::<()>::new(); #[derive(Debug)] @@ -1013,6 +1408,7 @@ where PeerMessage(RxJsonRpcMessage), ToSink(TxJsonRpcMessage), SendTaskResult(SendTaskResult), + ResponseSendTaskResult(Result<(), tokio::task::JoinError>), } let quit_reason = loop { @@ -1058,6 +1454,11 @@ where } } } + result = response_send_tasks.join_next(), if !response_send_tasks.is_empty() => { + Event::ResponseSendTaskResult( + result.expect("non-empty response send task set") + ) + } _ = serve_loop_ct.cancelled() => { tracing::info!("task cancelled"); break QuitReason::Cancelled @@ -1068,11 +1469,10 @@ where tracing::trace!(?evt, "new event"); match evt { Event::SendTaskResult(SendTaskResult::Request { id, result }) => { - if let Err(e) = result { - if let Some(responder) = local_responder_pool.remove(&id) { - let _ = responder.send(Err(ServiceError::TransportSend(e))); + if let Err(e) = result + && let Some(responder) = local_responder_pool.remove(&id) { + let _ = responder.responder.send(Err(ServiceError::TransportSend(e))); } - } } Event::SendTaskResult(SendTaskResult::Notification { responder, @@ -1085,15 +1485,18 @@ where Ok(()) }; let _ = responder.send(response); - if let Some(param) = cancellation_param { - if let Some(request_id) = ¶m.request_id { - if let Some(responder) = local_responder_pool.remove(request_id) { + if let Some(param) = cancellation_param + && let Some(request_id) = ¶m.request_id + && let Some(responder) = local_responder_pool.remove(request_id) { tracing::info!(id = %request_id, reason = param.reason, "cancelled"); - let _response_result = responder.send(Err(ServiceError::Cancelled { + let _response_result = responder.responder.send(Err(ServiceError::Cancelled { reason: param.reason.clone(), })); } - } + } + Event::ResponseSendTaskResult(result) => { + if let Err(error) = result { + tracing::error!(%error, "response send task failed"); } } // response and error @@ -1103,9 +1506,11 @@ where JsonRpcMessage::Error(error) => error.id.as_ref(), _ => None, } { - if let Some(ct) = local_ct_pool.remove(id) { - ct.cancel(); - } + let Some(ct) = local_ct_pool.remove(id) else { + tracing::debug!(%id, "dropping response for cancelled request"); + continue; + }; + ct.cancel(); let send = transport.send(m); let current_span = tracing::Span::current(); response_send_tasks.spawn(async move { @@ -1121,7 +1526,10 @@ where id, responder, }) => { - local_responder_pool.insert(id.clone(), responder); + local_responder_pool.insert(id.clone(), PendingResponse { + responder, + is_subscription: R::is_subscription_request(&request), + }); let send = transport.send(JsonRpcMessage::request(request, id.clone())); { let id = id.clone(); @@ -1159,6 +1567,24 @@ where .. })) => { tracing::debug!(%id, ?request, "received request"); + if let Err(error) = R::enforce_peer_request_association( + &request, + peer.peer_info().as_deref(), + peer_request_association(&request, &local_responder_pool), + ) { + tracing::warn!(%id, message = %error.message, "rejected peer request"); + // send directly: the sink proxy path would drop the + // error since the request was never registered in + // local_ct_pool + let send = transport.send(JsonRpcMessage::error(error, Some(id))); + let current_span = tracing::Span::current(); + response_send_tasks.spawn(async move { + if let Err(error) = send.await { + tracing::error!(%error, "fail to send rejection error"); + } + }.instrument(current_span)); + continue; + } { let service = shared_service.clone(); let sink = sink_proxy_tx.clone(); @@ -1166,7 +1592,7 @@ where let context_ct = request_ct.child_token(); local_ct_pool.insert(id.clone(), request_ct); let mut extensions = Extensions::new(); - let mut meta = Meta::new(); + let mut meta = RequestMetaObject::new(); // avoid clone // swap meta firstly, otherwise progress token will be lost std::mem::swap(&mut meta, request.get_meta_mut()); @@ -1179,9 +1605,10 @@ where extensions, }; let current_span = tracing::Span::current(); + let handler_id = id.clone(); spawn_service_task(async move { - let result = service - .handle_request(request, context) + let result = ORIGINATING_REQUEST + .scope(handler_id, service.handle_request(request, context)) .await; let response = match result { Ok(result) => { @@ -1202,26 +1629,82 @@ where .. })) => { tracing::info!(?notification, "received notification"); - // catch cancelled notification - let mut notification = match notification.try_into() { - Ok::(cancelled) => { - if let Some(request_id) = &cancelled.params.request_id { - if let Some(ct) = local_ct_pool.remove(request_id) { - tracing::info!(id = %request_id, reason = cancelled.params.reason, "cancelled"); + R::invalidate_response_cache(&peer, ¬ification).await; + let subscription_id = + if let Some(cancelled) = R::peer_cancelled_params(¬ification) { + let request_id = cancelled.request_id.clone(); + // Modern servers cancel listen requests; legacy peers cancel + // requests they originated, even when both directions share an ID. + if R::peer_cancels_subscriptions(&peer) { + if let Some(request_id) = request_id.as_ref() { + if local_responder_pool.get(request_id) + .is_some_and(|pending| pending.is_subscription) + && let Some(pending) = local_responder_pool.remove(request_id) + { + let _ = pending.responder.send(Err(ServiceError::Cancelled { + reason: cancelled.reason.clone(), + })); + } else { + tracing::debug!(%request_id, "ignoring cancellation of unknown subscription"); + continue; + } + } + request_id + } else { + if let Some(request_id) = request_id.as_ref() + && let Some(ct) = local_ct_pool.remove(request_id) + { + tracing::info!(id = %request_id, reason = cancelled.reason, "cancelled"); ct.cancel(); } + None + } + } else { + notification.get_meta().subscription_id() + }; + if let Some(subscription_id) = subscription_id + && let Some((sender, capacity)) = + peer.subscription_sender(&subscription_id) + { + match sender.try_send(notification) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(_)) => { + tracing::warn!( + id = %subscription_id, + capacity, + "subscription notification buffer full" + ); + if R::IS_CLIENT + && let Some(responder) = + local_responder_pool.remove(&subscription_id) + { + let _ = responder.responder + .send(Err(ServiceError::SubscriptionLagged { capacity })); + } + peer.unregister_subscription(&subscription_id); + peer.try_cancel_request( + subscription_id, + Some("subscription notification buffer full".to_owned()), + ); + } + Err(mpsc::error::TrySendError::Closed(_)) => { + peer.unregister_subscription(&subscription_id); + peer.try_cancel_request( + subscription_id, + Some("subscription notification receiver closed".to_owned()), + ); } - cancelled.into() } - Err(notification) => notification, - }; + continue; + } + let mut notification = notification; if let Some(progress_token) = notification.progress_token() { peer.notify_progress_timeout_watcher(progress_token).await; } { let service = shared_service.clone(); let mut extensions = Extensions::new(); - let mut meta = Meta::new(); + let mut meta = NotificationMetaObject::new(); // avoid clone std::mem::swap(&mut extensions, notification.extensions_mut()); std::mem::swap(&mut meta, notification.get_meta_mut()); @@ -1244,8 +1727,10 @@ where id, .. })) => { - if let Some(responder) = local_responder_pool.remove(&id) { - let response_result = responder.send(Ok(result)); + if let Some(responder) = + remove_pending_request(&mut local_responder_pool, &id) + { + let response_result = responder.responder.send(Ok(result)); if let Err(_error) = response_result { tracing::warn!(%id, "Error sending response"); } @@ -1258,8 +1743,15 @@ where tracing::debug!(?error, "received id-less peer error"); continue; }; - if let Some(responder) = local_responder_pool.remove(&id) { - let _response_result = responder.send(Err(ServiceError::McpError(error))); + if let Some(responder) = + remove_pending_request(&mut local_responder_pool, &id) + { + let service_error = if error.is_transport_closed() { + ServiceError::TransportClosed + } else { + ServiceError::McpError(error) + }; + let _response_result = responder.responder.send(Err(service_error)); if let Err(_error) = _response_result { tracing::warn!(%id, "Error sending response"); } @@ -1293,6 +1785,17 @@ where // Then drain any handler responses still in the channel // (handlers that finished after the loop broke). while let Some(m) = sink_proxy_rx.recv().await { + if let Some(id) = match &m { + JsonRpcMessage::Response(response) => Some(&response.id), + JsonRpcMessage::Error(error) => error.id.as_ref(), + _ => None, + } { + let Some(ct) = local_ct_pool.remove(id) else { + tracing::debug!(%id, "dropping response for cancelled request"); + continue; + }; + ct.cancel(); + } if let Err(error) = transport.send(m).await { tracing::error!(%error, "failed to send pending response during drain"); break; @@ -1320,3 +1823,119 @@ where dg: ct.drop_guard(), } } + +#[cfg(all(test, feature = "server"))] +mod sep2260_marker_tests { + use std::sync::Arc; + + use super::*; + use crate::model::{PingRequest, RequestId, ServerRequest}; + + fn ping() -> ServerRequest { + ServerRequest::PingRequest(PingRequest { + method: Default::default(), + extensions: Default::default(), + }) + } + + async fn send_and_capture(scope: Option) -> ::Req { + // peer_info None keeps enforcement non-strict; only the sink message matters. + let (peer, mut rx) = + Peer::::new(Arc::new(AtomicU32RequestIdProvider::default()), None); + let send = peer.send_request_with_option(ping(), PeerRequestOptions::no_options()); + let _handle = match scope { + Some(id) => ORIGINATING_REQUEST.scope(id, send).await.unwrap(), + None => send.await.unwrap(), + }; + let PeerSinkMessage::Request { request, .. } = rx.recv().await.expect("sink message") + else { + panic!("expected a request sink message"); + }; + request + } + + #[tokio::test] + async fn outbound_request_carries_originating_id_when_in_scope() { + let request = send_and_capture(Some(RequestId::Number(7))).await; + let marker = request + .extensions() + .get::() + .expect("marker attached"); + assert_eq!(marker.0, RequestId::Number(7)); + } + + #[tokio::test] + async fn outbound_request_has_no_marker_outside_scope() { + let request = send_and_capture(None).await; + assert!(request.extensions().get::().is_none()); + } + + #[test] + #[expect( + deprecated, + reason = "Sampling is deprecated by SEP-2577 but remains the canonical restricted request" + )] + fn peer_request_association_maps_stream_origin() { + use std::collections::HashMap; + + use crate::model::{ + CreateMessageRequest, CreateMessageRequestParams, SamplingMessage, ServerRequest, + }; + + fn sampling(origin: Option) -> ServerRequest { + let mut request = CreateMessageRequest::new(CreateMessageRequestParams::new( + vec![SamplingMessage::user_text("hi")], + 16, + )); + if let Some(origin) = origin { + request.extensions.insert(origin); + } + ServerRequest::CreateMessageRequest(request) + } + + let empty: HashMap = HashMap::new(); + let in_flight: HashMap = HashMap::from([(RequestId::Number(7), ())]); + + // No marker (stdio): coarse signal. + assert_eq!( + peer_request_association(&sampling(None), &in_flight), + PeerRequestAssociation::Unknown { + has_pending_outbound_request: true + } + ); + assert_eq!( + peer_request_association(&sampling(None), &empty), + PeerRequestAssociation::Unknown { + has_pending_outbound_request: false + } + ); + // Standalone GET stream: unassociated even with requests in flight. + assert_eq!( + peer_request_association( + &sampling(Some(InboundStreamOrigin::Unassociated)), + &in_flight + ), + PeerRequestAssociation::Unassociated + ); + // Originating POST stream of an in-flight request: associated. + assert_eq!( + peer_request_association( + &sampling(Some(InboundStreamOrigin::OutboundRequest( + RequestId::Number(7) + ))), + &in_flight + ), + PeerRequestAssociation::Associated + ); + // Stream of a request that is no longer in flight: unassociated. + assert_eq!( + peer_request_association( + &sampling(Some(InboundStreamOrigin::OutboundRequest( + RequestId::Number(8) + ))), + &in_flight + ), + PeerRequestAssociation::Unassociated + ); + } +} diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 05c2749..7d0d794 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -1,24 +1,35 @@ // Sampling/Roots/Logging are SEP-2577-deprecated; internal references are expected. #![expect(deprecated)] -use std::borrow::Cow; +pub(super) mod cache; +use std::{borrow::Cow, num::NonZeroUsize, sync::Arc, time::Duration}; + +use cache::CacheGeneration; +pub use cache::{ClientCacheConfig, MAX_CLIENT_CACHE_TTL}; use thiserror::Error; use super::*; use crate::{ model::{ - ArgumentInfo, CallToolRequest, CallToolRequestParams, CallToolResult, - CancelledNotification, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage, - ClientNotification, ClientRequest, ClientResult, CompleteRequest, CompleteRequestParams, - CompleteResult, CompletionContext, CompletionInfo, ErrorData, GetPromptRequest, - GetPromptRequestParams, GetPromptResult, InitializeRequest, InitializedNotification, - JsonRpcResponse, ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest, - ListResourceTemplatesResult, ListResourcesRequest, ListResourcesResult, ListToolsRequest, - ListToolsResult, PaginatedRequestParams, ProgressNotification, ProgressNotificationParam, - ReadResourceRequest, ReadResourceRequestParams, ReadResourceResult, Reference, RequestId, - RootsListChangedNotification, ServerInfo, ServerJsonRpcMessage, ServerNotification, - ServerRequest, ServerResult, SetLevelRequest, SetLevelRequestParams, SubscribeRequest, - SubscribeRequestParams, UnsubscribeRequest, UnsubscribeRequestParams, + ArgumentInfo, CacheScope, CallToolRequest, CallToolRequestParams, CallToolResponse, + CallToolResult, CancelTaskParams, CancelTaskRequest, CancelledNotification, + CancelledNotificationParam, ClientConfig, ClientJsonRpcMessage, ClientNotification, + ClientRequest, ClientResult, CompleteRequest, CompleteRequestParams, CompleteResult, + CompletionContext, CompletionInfo, DEFAULT_MRTR_MAX_ROUNDS, DiscoverRequest, + DiscoverRequestParams, DiscoverResult, ErrorData, GetExtensions, GetMeta, GetPromptRequest, + GetPromptRequestParams, GetPromptResponse, GetPromptResult, GetTaskParams, GetTaskRequest, + GetTaskResult, InitializeRequest, InitializeRequestParams, InitializedNotification, + InputRequest, InputRequiredResult, InputResponses, JsonRpcResponse, ListPromptsRequest, + ListPromptsResult, ListResourceTemplatesRequest, ListResourceTemplatesResult, + ListResourcesRequest, ListResourcesResult, ListToolsRequest, ListToolsResult, + NumberOrString, PaginatedRequestParams, ProgressNotification, ProgressNotificationParam, + ProtocolVersion, ReadResourceRequest, ReadResourceRequestParams, ReadResourceResponse, + ReadResourceResult, Reference, RequestId, RequestMetaObject, RootsListChangedNotification, + ServerJsonRpcMessage, ServerNotification, ServerPeerInfo, ServerRequest, ServerResult, + SetLevelRequest, SetLevelRequestParams, SubscribeRequest, SubscribeRequestParams, + SubscriptionFilter, SubscriptionsListenRequest, SubscriptionsListenRequestParams, + SubscriptionsListenResult, UnsubscribeRequest, UnsubscribeRequestParams, UpdateTaskParams, + UpdateTaskRequest, }, transport::DynamicTransportError, }; @@ -38,6 +49,14 @@ pub enum ClientInitializeError { #[error("conflict initialized response id: expected {0}, got {1}")] ConflictInitResponseId(RequestId, RequestId), + #[error( + "uncorrelated error response: expected id {expected}, error response carried {received}" + )] + UncorrelatedErrorResponse { + expected: RequestId, + received: RequestId, + }, + #[error("connection closed: {0}")] ConnectionClosed(String), @@ -50,8 +69,25 @@ pub enum ClientInitializeError { #[error("JSON-RPC error: {0}")] JsonRpcError(ErrorData), + #[error( + "no compatible protocol version (client: {client_supported:?}, server: {server_supported:?})" + )] + NoCompatibleProtocolVersion { + client_supported: Vec, + server_supported: Vec, + }, + + #[error("discover startup requires at least one preferred protocol version")] + NoPreferredProtocolVersion, + #[error("Cancelled")] Cancelled, + + #[error("discover and legacy initialize both failed")] + LegacyFallbackFailed { + discover: Box, + fallback: Box, + }, } impl ClientInitializeError { @@ -64,6 +100,48 @@ impl ClientInitializeError { context: context.into(), } } + + /// The `WWW-Authenticate` challenge from the 401/403 the transport hit + /// during initialization, if that is why initialization failed. + /// + /// This is the trigger of the reactive OAuth flow: feed the challenge to + /// `AuthorizationRequest::with_challenge` to authorize, then reconnect. + #[cfg(feature = "transport-streamable-http-client")] + pub fn auth_challenge(&self) -> Option<&str> { + use crate::transport::streamable_http_client::{AuthRequiredError, InsufficientScopeError}; + + let error = match self { + Self::TransportError { error, .. } => error, + // A 401/403 in the fallback phase is still actionable. + Self::LegacyFallbackFailed { fallback, .. } => { + return fallback.auth_challenge(); + } + _ => return None, + }; + let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error.error.as_ref()); + while let Some(current) = source { + if let Some(auth_required) = current.downcast_ref::() { + return Some(&auth_required.www_authenticate_header); + } + if let Some(insufficient_scope) = current.downcast_ref::() { + return Some(&insufficient_scope.www_authenticate_header); + } + source = current.source(); + } + None + } + + /// Returns whether client initialization failed because authorization is required. + /// + /// This covers both missing or expired local OAuth authorization and an HTTP + /// authorization challenge from the MCP server. + pub fn is_authorization_required(&self) -> bool { + match self { + Self::TransportError { error, .. } => error.is_authorization_required(), + Self::LegacyFallbackFailed { fallback, .. } => fallback.is_authorization_required(), + _ => false, + } + } } /// Helper function to get the next message from the stream @@ -80,13 +158,20 @@ where .ok_or_else(|| ClientInitializeError::ConnectionClosed(context.to_string())) } -/// Helper function to expect a response from the stream +/// Helper function to expect a response from the stream, correlated to +/// `expected_id`. +/// +/// Both success and error responses are checked here: a mismatched id on a +/// success response is `ConflictInitResponseId`; on an error response (whose +/// `id` is optional per spec) it is `UncorrelatedErrorResponse`. The caller +/// never sees an uncorrelated response. async fn expect_response( transport: &mut T, context: &str, service: &S, peer: Peer, -) -> Result<(ServerResult, RequestId), ClientInitializeError> + expected_id: &RequestId, +) -> Result where T: Transport, S: Service, @@ -94,13 +179,29 @@ where loop { let message = expect_next_message(transport, context).await?; match message { - // Expected message to complete the initialization ServerJsonRpcMessage::Response(JsonRpcResponse { id, result, .. }) => { - break Ok((result, id)); + if !expected_id.matches_response_id(&id) { + return Err(ClientInitializeError::ConflictInitResponseId( + expected_id.clone(), + id, + )); + } + return Ok(result); } - // Handle JSON-RPC error responses ServerJsonRpcMessage::Error(error) => { - break Err(ClientInitializeError::JsonRpcError(error.error)); + return Err(match &error.id { + Some(id) if expected_id.matches_response_id(id) => { + ClientInitializeError::JsonRpcError(error.error) + } + // Spec: error id is optional; a server that cannot read + // the request id omits it. The error is still a response + // to our request, so it remains available to the caller. + None => ClientInitializeError::JsonRpcError(error.error), + Some(id) => ClientInitializeError::UncorrelatedErrorResponse { + expected: expected_id.clone(), + received: id.clone(), + }, + }); } // Server could send logging messages before handshake ServerJsonRpcMessage::Notification(mut notification) => { @@ -113,11 +214,11 @@ where let mut context = NotificationContext { peer: peer.clone(), - meta: Meta::default(), + meta: NotificationMetaObject::default(), extensions: Extensions::default(), }; - if let Some(meta) = logging.extensions.get_mut::() { + if let Some(meta) = logging.extensions.get_mut::() { std::mem::swap(&mut context.meta, meta); } std::mem::swap(&mut context.extensions, &mut logging.extensions); @@ -145,6 +246,19 @@ where #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] pub struct RoleClient; +/// Select the first client-preferred protocol version supported by the server. +/// +/// Returns `None` when no version is shared. +pub fn select_protocol_version( + client_preference: &[ProtocolVersion], + server_supported: &[ProtocolVersion], +) -> Option { + client_preference + .iter() + .find(|version| server_supported.contains(version)) + .cloned() +} + impl ServiceRole for RoleClient { type Req = ClientRequest; type Resp = ClientResult; @@ -152,14 +266,414 @@ impl ServiceRole for RoleClient { type PeerReq = ServerRequest; type PeerResp = ServerResult; type PeerNot = ServerNotification; - type Info = ClientInfo; - type PeerInfo = ServerInfo; + type Info = ClientConfig; + type PeerInfo = ServerPeerInfo; type InitializeError = ClientInitializeError; const IS_CLIENT: bool = true; + + fn configure_direct_peer(peer: &Peer, info: &Self::Info) { + let Some(server_info) = peer.peer_info() else { + return; + }; + if server_info.protocol_version.as_str() < ProtocolVersion::V_2026_07_28.as_str() { + return; + } + peer.set_client_request_metadata(ClientRequestMetadata { + protocol_version: server_info.protocol_version.clone(), + client_info: info.client_info.clone(), + client_capabilities: info.capabilities.clone(), + }); + } + + fn peer_cancelled_params(notification: &Self::PeerNot) -> Option<&CancelledNotificationParam> { + match notification { + ServerNotification::CancelledNotification(notification) => Some(¬ification.params), + _ => None, + } + } + + fn peer_cancels_subscriptions(peer: &Peer) -> bool { + // Discovery keeps modern lifecycle semantics even with an older application version. + !super::uses_legacy_lifecycle( + peer.peer_info() + .as_deref() + .map(|info| &info.protocol_version), + peer.client_request_metadata.get().is_some(), + ) + } + + fn is_subscription_request(request: &Self::Req) -> bool { + matches!(request, ClientRequest::SubscriptionsListenRequest(_)) + } + + // SEP-2260: reject restricted server requests that arrived unassociated + // with any in-flight outbound request. Without stream separation + // (`Unknown`) the coarse in-flight check under-approximates the SHOULD. + fn enforce_peer_request_association( + peer_request: &Self::PeerReq, + peer_info: Option<&Self::PeerInfo>, + association: PeerRequestAssociation, + ) -> Result<(), ErrorData> { + let restricted = matches!( + peer_request, + ServerRequest::CreateMessageRequest(_) + | ServerRequest::ListRootsRequest(_) + | ServerRequest::ElicitRequest(_) + ); + if !restricted { + return Ok(()); + } + let strict = + peer_info.is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28); + if !strict { + return Ok(()); + } + let associated = match association { + PeerRequestAssociation::Associated => true, + PeerRequestAssociation::Unassociated => false, + PeerRequestAssociation::Unknown { + has_pending_outbound_request, + } => has_pending_outbound_request, + }; + if associated { + Ok(()) + } else { + Err(ErrorData::invalid_params( + "SEP-2260: server-to-client requests must be associated with an in-flight client request", + None, + )) + } + } + + async fn invalidate_response_cache(peer: &Peer, notification: &Self::PeerNot) { + match notification { + ServerNotification::ResourceUpdatedNotification(notification) => { + peer.invalidate_resource_read_cache(¬ification.params.uri) + .await; + } + ServerNotification::ResourceListChangedNotification(_) => { + peer.invalidate_resource_list_cache().await; + } + ServerNotification::ToolListChangedNotification(_) => { + peer.invalidate_tool_cache().await; + } + ServerNotification::PromptListChangedNotification(_) => { + peer.invalidate_prompt_cache().await; + } + _ => {} + } + } } pub type ServerSink = Peer; +/// Default number of notifications buffered for one subscription. +pub const DEFAULT_SUBSCRIPTION_CHANNEL_CAPACITY: usize = 64; + +/// How a client-side subscription stream ended. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub enum SubscriptionEnd { + /// The server returned a final `SubscriptionsListenResult`. + Graceful(SubscriptionsListenResult), + /// The transport closed without a final result. Call `Peer::listen` again + /// after reconnecting; subscription streams are not resumable. + Abrupt, + /// The subscription was explicitly cancelled by either peer. + Cancelled, + /// The consumer did not drain notifications before the channel filled. + Lagged { capacity: usize }, +} + +/// Handle for one active `subscriptions/listen` request. +#[derive(Debug)] +pub struct Subscription { + id: RequestId, + acknowledged: SubscriptionFilter, + notifications: tokio::sync::mpsc::Receiver, + request: Option>, + end: Option, +} + +type SubscriptionResponse = + Result, tokio::sync::oneshot::error::RecvError>; + +struct PendingSubscriptionRequest { + handle: Option>, +} + +impl PendingSubscriptionRequest { + fn new(handle: RequestHandle) -> Self { + Self { + handle: Some(handle), + } + } + + async fn recv(&mut self) -> Option { + let handle = self.handle.as_mut()?; + Some((&mut handle.rx).await) + } + + fn take(&mut self) -> Option> { + self.handle.take() + } + + fn unregister(&self, id: &RequestId) { + if let Some(handle) = self.handle.as_ref() { + handle.peer.unregister_subscription(id); + } + } + + fn disarm(&mut self) { + self.handle.take(); + } + + async fn cancel(&mut self, reason: &'static str) { + if let Some(handle) = self.handle.take() { + let _ = handle.cancel(Some(reason.to_owned())).await; + } + } +} + +impl Drop for PendingSubscriptionRequest { + fn drop(&mut self) { + let Some(handle) = self.handle.take() else { + return; + }; + handle.peer.unregister_subscription(&handle.id); + handle.peer.try_cancel_request( + handle.id, + Some("subscription establishment cancelled".to_owned()), + ); + } +} + +impl Subscription { + /// Return the originating listen request ID. + pub fn id(&self) -> &RequestId { + &self.id + } + + /// Return the notification filter accepted by the server. + pub fn acknowledged(&self) -> &SubscriptionFilter { + &self.acknowledged + } + + /// Return the terminal state after this subscription has ended. + pub fn end(&self) -> Option<&SubscriptionEnd> { + self.end.as_ref() + } + + /// Receive the next notification, or `None` after the subscription ends. + /// + /// A graceful final result and an abrupt transport close are distinguished + /// through [`Self::end`]. + /// + /// # Errors + /// + /// Returns a service or protocol error when the stream carries an invalid + /// message, an unexpected final result, or another request failure. + pub async fn next(&mut self) -> Result, ServiceError> { + if self.end.is_some() { + return Ok(None); + } + let Some(request) = self.request.as_mut() else { + self.end = Some(SubscriptionEnd::Abrupt); + return Ok(None); + }; + + tokio::select! { + biased; + notification = self.notifications.recv() => { + let Some(notification) = notification else { + let response = (&mut request.rx).await; + return self.handle_response(response); + }; + if let ServerNotification::CancelledNotification(cancelled) = ¬ification { + if cancelled.params.request_id.as_ref() != Some(&self.id) { + self.cancel_as_abrupt("subscription cancellation ID mismatch") + .await; + return Err(ServiceError::UnexpectedResponse); + } + self.finish(SubscriptionEnd::Cancelled); + return Ok(None); + } + if notification.get_meta().subscription_id().as_ref() != Some(&self.id) { + self.cancel_as_abrupt("subscription notification ID mismatch") + .await; + return Err(ServiceError::UnexpectedResponse); + } + if !self.accepts(¬ification) { + self.cancel_as_abrupt( + "subscription notification was outside the acknowledged filter", + ) + .await; + return Err(ServiceError::UnexpectedResponse); + } + Ok(Some(notification)) + } + response = &mut request.rx => { + self.handle_response(response) + } + } + } + + /// Cancel this subscription. + /// + /// # Errors + /// + /// Returns a transport error when the cancellation signal cannot be sent. + pub async fn cancel(&mut self) -> Result<(), ServiceError> { + self.cancel_with_reason(None).await + } + + /// Cancel this subscription with a diagnostic reason. + /// + /// # Errors + /// + /// Returns a transport error when the cancellation signal cannot be sent. + pub async fn cancel_with_reason(&mut self, reason: Option) -> Result<(), ServiceError> { + let Some(request) = self.request.take() else { + return Ok(()); + }; + request.cancel(reason).await?; + self.end = Some(SubscriptionEnd::Cancelled); + Ok(()) + } + + fn finish(&mut self, end: SubscriptionEnd) { + if let Some(request) = self.request.take() { + request.peer.unregister_subscription(&self.id); + } + self.end = Some(end); + } + + async fn cancel_as_abrupt(&mut self, reason: &'static str) { + if let Some(request) = self.request.take() { + let _ = request.cancel(Some(reason.to_owned())).await; + } + self.end = Some(SubscriptionEnd::Abrupt); + } + + fn accepts(&self, notification: &ServerNotification) -> bool { + match notification { + ServerNotification::ToolListChangedNotification(_) => { + self.acknowledged.tools_list_changed == Some(true) + } + ServerNotification::PromptListChangedNotification(_) => { + self.acknowledged.prompts_list_changed == Some(true) + } + ServerNotification::ResourceListChangedNotification(_) => { + self.acknowledged.resources_list_changed == Some(true) + } + ServerNotification::ResourceUpdatedNotification(update) => self + .acknowledged + .resource_subscriptions + .as_ref() + .is_some_and(|uris| uris.contains(&update.params.uri)), + ServerNotification::SubscriptionsAcknowledgedNotification(_) + | ServerNotification::CancelledNotification(_) + | ServerNotification::ProgressNotification(_) + | ServerNotification::LoggingMessageNotification(_) + | ServerNotification::TaskStatusNotification(_) + | ServerNotification::CustomNotification(_) => false, + } + } + + fn handle_response( + &mut self, + response: SubscriptionResponse, + ) -> Result, ServiceError> { + let response = match response { + Ok(response) => response, + Err(_) => { + self.finish(SubscriptionEnd::Abrupt); + return Ok(None); + } + }; + let response = match response { + Ok(response) => response, + Err(ServiceError::TransportClosed) => { + self.finish(SubscriptionEnd::Abrupt); + return Ok(None); + } + Err(ServiceError::SubscriptionLagged { capacity }) => { + self.finish(SubscriptionEnd::Lagged { capacity }); + return Ok(None); + } + Err(error) => { + self.finish(SubscriptionEnd::Abrupt); + return Err(error); + } + }; + let ServerResult::SubscriptionsListenResult(result) = response else { + self.finish(SubscriptionEnd::Abrupt); + return Err(ServiceError::UnexpectedResponse); + }; + if !result.result_type.is_complete() + || result.meta.subscription_id().as_ref() != Some(&self.id) + { + self.finish(SubscriptionEnd::Abrupt); + return Err(ServiceError::UnexpectedResponse); + } + self.finish(SubscriptionEnd::Graceful(result)); + Ok(None) + } +} + +impl Drop for Subscription { + fn drop(&mut self) { + let Some(request) = self.request.take() else { + return; + }; + request.peer.unregister_subscription(&self.id); + request.peer.try_cancel_request( + self.id.clone(), + Some("subscription handle dropped".to_owned()), + ); + } +} + +/// Selects how a client establishes its MCP lifecycle. +/// +/// Existing [`ServiceExt::serve`] behavior remains legacy initialization. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum ClientLifecycleMode { + /// Use the legacy `initialize` / `notifications/initialized` handshake. + Initialize, + /// Use `server/discover` and send self-contained per-request metadata. + Discover { + preferred_versions: Vec, + }, + /// Probe with `server/discover`, falling back when the peer reports that it is legacy or does + /// not respond within 10 seconds. + Auto { + preferred_versions: Vec, + legacy_version: Option, + }, +} + +const DEFAULT_AUTO_DISCOVER_TIMEOUT: Duration = Duration::from_secs(10); + +/// Client-specific lifecycle entry points. +pub trait ClientServiceExt: Service + Sized { + fn serve_with_lifecycle( + self, + transport: T, + lifecycle: ClientLifecycleMode, + ) -> impl Future, ClientInitializeError>> + + MaybeSendFuture + where + T: IntoTransport, + E: std::error::Error + Send + Sync + 'static, + { + serve_client_with_lifecycle(self, transport, lifecycle) + } +} + +impl> ClientServiceExt for S {} + impl> ServiceExt for S { fn serve_with_ct( self, @@ -185,7 +699,13 @@ where T: IntoTransport, E: std::error::Error + Send + Sync + 'static, { - serve_client_with_ct(service, transport, Default::default()).await + serve_client_with_lifecycle_and_ct( + service, + transport, + ClientLifecycleMode::Initialize, + Default::default(), + ) + .await } pub async fn serve_client_with_ct( @@ -193,13 +713,47 @@ pub async fn serve_client_with_ct( transport: T, ct: CancellationToken, ) -> Result, ClientInitializeError> +where + S: Service, + T: IntoTransport, + E: std::error::Error + Send + Sync + 'static, +{ + serve_client_with_lifecycle_and_ct(service, transport, ClientLifecycleMode::Initialize, ct) + .await +} + +pub async fn serve_client_with_lifecycle( + service: S, + transport: T, + lifecycle: ClientLifecycleMode, +) -> Result, ClientInitializeError> +where + S: Service, + T: IntoTransport, + E: std::error::Error + Send + Sync + 'static, +{ + serve_client_with_lifecycle_and_ct(service, transport, lifecycle, Default::default()).await +} + +pub async fn serve_client_with_lifecycle_and_ct( + service: S, + transport: T, + lifecycle: ClientLifecycleMode, + ct: CancellationToken, +) -> Result, ClientInitializeError> where S: Service, T: IntoTransport, E: std::error::Error + Send + Sync + 'static, { tokio::select! { - result = serve_client_with_ct_inner(service, transport.into_transport(), ct.clone()) => { result } + result = serve_client_with_ct_inner( + service, + transport.into_transport(), + lifecycle, + ct.clone(), + DEFAULT_AUTO_DISCOVER_TIMEOUT, + ) => { result } _ = ct.cancelled() => { Err(ClientInitializeError::Cancelled) } @@ -209,7 +763,9 @@ where async fn serve_client_with_ct_inner( service: S, transport: T, + lifecycle: ClientLifecycleMode, ct: CancellationToken, + auto_discover_timeout: Duration, ) -> Result, ClientInitializeError> where S: Service, @@ -217,12 +773,134 @@ where { let mut transport = transport.into_transport(); let id_provider = >::default(); + let (peer, peer_rx) = Peer::new(id_provider.clone(), None); + let client_info = service.get_info(); - // service + match lifecycle { + ClientLifecycleMode::Initialize => { + legacy_startup(&service, &mut transport, &id_provider, &peer, client_info).await?; + } + ClientLifecycleMode::Discover { preferred_versions } => { + match discover_startup( + &service, + &mut transport, + &id_provider, + &peer, + &client_info, + preferred_versions, + ) + .await? + { + DiscoverOutcome::Modern => {} + // Discover mode does not fall back; a legacy server is an error. + DiscoverOutcome::Legacy(error) => return Err(*error), + } + } + ClientLifecycleMode::Auto { + preferred_versions, + legacy_version, + } => { + let discover_result = tokio::time::timeout( + auto_discover_timeout, + discover_startup( + &service, + &mut transport, + &id_provider, + &peer, + &client_info, + preferred_versions, + ), + ) + .await; + match discover_result { + Ok(Ok(DiscoverOutcome::Modern)) => {} + Ok(Ok(DiscoverOutcome::Legacy(discover_error))) => { + let mut legacy_info = client_info; + if let Some(version) = legacy_version { + legacy_info.protocol_version = version; + } + if let Err(fallback_error) = + legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info) + .await + { + return Err(ClientInitializeError::LegacyFallbackFailed { + discover: discover_error, + fallback: Box::new(fallback_error), + }); + } + } + Ok(Err(error)) => return Err(error), + Err(_) => { + let mut legacy_info = client_info; + if let Some(version) = legacy_version { + legacy_info.protocol_version = version; + } + legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info) + .await?; + } + } + } + } + Ok(serve_inner( + service, + transport, + peer, + peer_rx, + VecDeque::new(), + ct, + )) +} + +/// Modern-era JSON-RPC error codes a server can return from `server/discover` +/// without being legacy. Version negotiation (`UNSUPPORTED_PROTOCOL_VERSION`) +/// is handled by `discover_startup`'s own retry loop and never reaches the +/// classification below. +/// +/// `ErrorCode` is an open integer type, so this cannot be exhaustive: if a +/// future revision adds another modern-era rejection code, add it here. +fn is_modern_rejection_code(code: crate::model::ErrorCode) -> bool { + matches!( + code, + crate::model::ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY + | crate::model::ErrorCode::HEADER_MISMATCH + ) +} + +/// The outcome of a `server/discover` probe, classified at the point where all +/// the context (request id, response correlation, transport state) is still +/// available. +/// +/// `Legacy` is returned only when the probe produced a complete, correlated +/// JSON-RPC error whose code is not a modern-era rejection — i.e. the +/// transport is in a known-good state and the error identifies the peer as +/// legacy per the 2026-07-28 backward-compatibility guidance. Every other +/// failure (transport error, uncorrelated response, modern rejection, etc.) +/// becomes `Err` so the caller surfaces it instead of retrying. +enum DiscoverOutcome { + /// The server speaks the modern protocol; discovery succeeded. + Modern, + /// The server is legacy: discovery received a correlated, non-modern + /// JSON-RPC error. The transport is still usable for a legacy `initialize` + /// handshake. The original error is preserved so a failed fallback can + /// report both phases. + Legacy(Box), +} + +async fn legacy_startup( + service: &S, + transport: &mut T, + id_provider: &Arc, + peer: &Peer, + client_info: InitializeRequestParams, +) -> Result<(), ClientInitializeError> +where + S: Service, + T: Transport + 'static, +{ let id = id_provider.next_request_id(); let init_request = InitializeRequest { method: Default::default(), - params: service.get_info(), + params: client_info, extensions: Default::default(), }; transport @@ -236,27 +914,13 @@ where context: "send initialize request".into(), })?; - let (peer, peer_rx) = Peer::new(id_provider, None); - - let (response, response_id) = expect_response( - &mut transport, - "initialize response", - &service, - peer.clone(), - ) - .await?; - - if id != response_id { - return Err(ClientInitializeError::ConflictInitResponseId( - id, - response_id, - )); - } + let response = + expect_response(transport, "initialize response", service, peer.clone(), &id).await?; let ServerResult::InitializeResult(initialize_result) = response else { return Err(ClientInitializeError::ExpectedInitResult(Some(response))); }; - peer.set_peer_info(initialize_result); + peer.set_peer_info(initialize_result.into()); // send notification let notification = ClientJsonRpcMessage::notification( @@ -268,7 +932,165 @@ where transport.send(notification).await.map_err(|error| { ClientInitializeError::transport::(error, "send initialized notification") })?; - Ok(serve_inner(service, transport, peer, peer_rx, ct)) + Ok(()) +} + +async fn discover_startup( + service: &S, + transport: &mut T, + id_provider: &Arc, + peer: &Peer, + client_info: &InitializeRequestParams, + preferred_versions: Vec, +) -> Result +where + S: Service, + T: Transport + 'static, +{ + if preferred_versions.is_empty() { + return Err(ClientInitializeError::NoPreferredProtocolVersion); + } + + let mut attempted = Vec::new(); + let mut candidate = preferred_versions[0].clone(); + loop { + attempted.push(candidate.clone()); + + let meta = RequestMetaObject::with_client_context( + candidate.clone(), + client_info.client_info.clone(), + client_info.capabilities.clone(), + ); + let mut discover = DiscoverRequest::new(DiscoverRequestParams {}); + discover.extensions.insert(meta); + let id = id_provider.next_request_id(); + transport + .send(ClientJsonRpcMessage::request( + ClientRequest::DiscoverRequest(discover), + id.clone(), + )) + .await + .map_err(|error| { + ClientInitializeError::transport::(error, "send discover request") + })?; + + match expect_response(transport, "discover response", service, peer.clone(), &id).await { + Ok(ServerResult::DiscoverResult(result)) => { + let Some(selected) = + select_protocol_version(&preferred_versions, &result.supported_versions) + else { + return Err(ClientInitializeError::NoCompatibleProtocolVersion { + client_supported: preferred_versions, + server_supported: result.supported_versions, + }); + }; + peer.set_peer_info(ServerPeerInfo::from_discover_result( + selected.clone(), + result, + )); + peer.set_client_request_metadata(ClientRequestMetadata { + protocol_version: selected, + client_info: client_info.client_info.clone(), + client_capabilities: client_info.capabilities.clone(), + }); + return Ok(DiscoverOutcome::Modern); + } + Ok(response) => { + return Err(ClientInitializeError::ExpectedInitResult(Some(response))); + } + Err(ClientInitializeError::JsonRpcError(error)) + if error.code == crate::model::ErrorCode::UNSUPPORTED_PROTOCOL_VERSION => + { + let supported = error + .data + .as_ref() + .and_then(|data| data.get("supported")) + .cloned() + .and_then(|value| serde_json::from_value::>(value).ok()) + .unwrap_or_default(); + let may_retry_current = attempted + .iter() + .filter(|version| *version == &candidate) + .count() + == 1; + let next = preferred_versions + .iter() + .find(|version| { + supported.contains(version) + && (!attempted.contains(version) + || (may_retry_current && *version == &candidate)) + }) + .cloned(); + let Some(next) = next else { + return Err(ClientInitializeError::NoCompatibleProtocolVersion { + client_supported: preferred_versions, + server_supported: supported, + }); + }; + candidate = next; + } + // A correlated JSON-RPC error that is not a modern-era rejection + // and not a version-negotiation signal: the server is legacy. + // The transport delivered a complete response, so a legacy + // `initialize` can follow on the same connection. + Err(error) + if matches!( + &error, + ClientInitializeError::JsonRpcError(data) + if !is_modern_rejection_code(data.code) + ) => + { + return Ok(DiscoverOutcome::Legacy(Box::new(error))); + } + Err(error) => return Err(error), + } + } +} + +const DISCOVER_CACHE_PREFIX: &str = "server/discover:"; +const TOOL_LIST_CACHE_PREFIX: &str = "tools/list:"; +const PROMPT_LIST_CACHE_PREFIX: &str = "prompts/list:"; +const RESOURCE_LIST_CACHE_PREFIX: &str = "resources/list:"; +const RESOURCE_TEMPLATE_LIST_CACHE_PREFIX: &str = "resources/templates/list:"; +const RESOURCE_READ_CACHE_PREFIX: &str = "resources/read:"; + +// Cache keys are built only from the request method plus the parameters that +// affect the result (SEP-2549). Request `_meta` (progress tokens, trace +// context, etc.) does not affect the result, so it is deliberately excluded to +// avoid fragmenting the cache across otherwise-identical requests. +fn discover_cache_key() -> String { + // `server/discover` carries no result-affecting parameters. + DISCOVER_CACHE_PREFIX.to_string() +} + +fn list_response_cache_key(prefix: &str, params: &Option) -> String { + // Only the pagination cursor affects which page is returned. + let cursor = params.as_ref().and_then(|params| params.cursor.as_deref()); + let cursor = + serde_json::to_string(&cursor).expect("serializing a pagination cursor cannot fail"); + format!("{prefix}{cursor}") +} + +fn resource_read_cache_key(params: &ReadResourceRequestParams) -> Option { + // MRTR retries depend on inputs that are not part of the cache key and MUST + // NOT be cached. + if params.input_responses.is_some() || params.request_state.is_some() { + return None; + } + // Only the URI affects the result. + Some(resource_read_cache_prefix_for_uri(¶ms.uri)) +} + +fn resource_read_cache_prefix_for_uri(uri: &str) -> String { + let uri = serde_json::to_string(uri).expect("serializing a resource URI cannot fail"); + format!("{RESOURCE_READ_CACHE_PREFIX}{uri}:") +} + +fn request_uses_cursor(params: &Option) -> bool { + params + .as_ref() + .and_then(|params| params.cursor.as_ref()) + .is_some() } macro_rules! method { @@ -361,6 +1183,322 @@ macro_rules! method { } impl Peer { + /// Open a long-lived notification subscription and wait for its acknowledgment. + /// + /// Notifications routed to the returned [`Subscription`] are not also + /// delivered through [`ClientHandler`](crate::ClientHandler) callbacks. + /// + /// # Errors + /// + /// Returns a service, transport, or protocol error when the request cannot + /// be established or the acknowledgment is invalid. + pub async fn listen( + &self, + notifications: SubscriptionFilter, + ) -> Result { + self.listen_with_channel_capacity_inner( + notifications, + DEFAULT_SUBSCRIPTION_CHANNEL_CAPACITY, + ) + .await + } + + /// Open a subscription with an explicit notification buffer capacity. + /// + /// Notifications routed to the returned [`Subscription`] are not also + /// delivered through [`ClientHandler`](crate::ClientHandler) callbacks. + /// + /// # Errors + /// + /// Returns a service, transport, or protocol error when the request cannot + /// be established or the acknowledgment is invalid. + pub async fn listen_with_capacity( + &self, + notifications: SubscriptionFilter, + channel_capacity: NonZeroUsize, + ) -> Result { + self.listen_with_channel_capacity_inner(notifications, channel_capacity.get()) + .await + } + + async fn listen_with_channel_capacity_inner( + &self, + notifications: SubscriptionFilter, + channel_capacity: usize, + ) -> Result { + let request = ClientRequest::SubscriptionsListenRequest(SubscriptionsListenRequest::new( + SubscriptionsListenRequestParams::new(notifications.clone()), + )); + let (handle, mut subscription_notifications) = self + .send_subscription_request(request, PeerRequestOptions::no_options(), channel_capacity) + .await?; + let id = handle.id.clone(); + let mut pending = PendingSubscriptionRequest::new(handle); + + tokio::select! { + biased; + notification = subscription_notifications.recv() => { + let Some(notification) = notification else { + pending.cancel("subscription stream closed before acknowledgment").await; + return Err(ServiceError::TransportClosed); + }; + if notification.get_meta().subscription_id().as_ref() != Some(&id) { + pending.cancel("subscription acknowledgment ID mismatch").await; + return Err(ServiceError::UnexpectedResponse); + } + let ServerNotification::SubscriptionsAcknowledgedNotification( + acknowledgment, + ) = notification else { + pending.cancel("notification received before subscription acknowledgment").await; + return Err(ServiceError::UnexpectedResponse); + }; + let accepted = acknowledgment.params.notifications; + if !accepted.is_subset_of(¬ifications) { + pending.cancel("subscription acknowledged an unrequested filter").await; + return Err(ServiceError::UnexpectedResponse); + } + let Some(handle) = pending.take() else { + return Err(ServiceError::TransportClosed); + }; + Ok(Subscription { + id, + acknowledged: accepted, + notifications: subscription_notifications, + request: Some(handle), + end: None, + }) + } + response = pending.recv() => { + pending.unregister(&id); + pending.disarm(); + let Some(response) = response else { + return Err(ServiceError::TransportClosed); + }; + match response { + Ok(Err(error)) => Err(error), + Ok(Ok(_)) => Err(ServiceError::UnexpectedResponse), + Err(_) => Err(ServiceError::TransportClosed), + } + } + } + } + + /// Discover the server's supported protocol versions and capabilities. + /// + /// The high-level client currently exposes this peer only after initialization; + /// pre-initialization probing is planned as follow-up work. + pub async fn discover(&self, meta: RequestMetaObject) -> Result { + let cache_key = discover_cache_key(); + if let Some(ServerResult::DiscoverResult(result)) = self.cached_response(&cache_key).await { + return Ok(result); + } + let generation = self.capture_response_cache_generation().await; + let mut request = DiscoverRequest::new(DiscoverRequestParams {}); + request.extensions.insert(meta); + let result = self + .send_request(ClientRequest::DiscoverRequest(request)) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if let Some(ServerResult::DiscoverResult(result)) = + self.stale_cached_response(&cache_key).await + { + return Ok(result); + } + return Err(error); + } + }; + match result { + ServerResult::DiscoverResult(result) => { + self.cache_result( + Some(cache_key), + Some(result.ttl_ms), + Some(result.cache_scope), + generation, + ServerResult::DiscoverResult(result.clone()), + ) + .await; + Ok(result) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } + + async fn cache_result( + &self, + cache_key: Option, + ttl_ms: Option, + cache_scope: Option, + generation: CacheGeneration, + result: ServerResult, + ) { + let Some(cache_key) = cache_key else { + return; + }; + self.cache_response_with_generation(cache_key, result, ttl_ms, cache_scope, generation) + .await; + } + + pub(crate) async fn invalidate_tool_cache(&self) { + self.invalidate_cached_responses(TOOL_LIST_CACHE_PREFIX) + .await; + } + + pub(crate) async fn invalidate_prompt_cache(&self) { + self.invalidate_cached_responses(PROMPT_LIST_CACHE_PREFIX) + .await; + } + + pub(crate) async fn invalidate_resource_list_cache(&self) { + self.invalidate_cached_responses(RESOURCE_LIST_CACHE_PREFIX) + .await; + self.invalidate_cached_responses(RESOURCE_TEMPLATE_LIST_CACHE_PREFIX) + .await; + } + + pub(crate) async fn invalidate_resource_read_cache(&self, uri: &str) { + self.invalidate_cached_responses(&resource_read_cache_prefix_for_uri(uri)) + .await; + } + + /// Send one `tools/call` request and return either a final result or an MRTR + /// `InputRequiredResult` without driving any follow-up rounds. + pub async fn call_tool_once( + &self, + params: CallToolRequestParams, + ) -> Result { + let result = self + .send_request(ClientRequest::CallToolRequest(CallToolRequest { + method: Default::default(), + params, + extensions: Default::default(), + })) + .await?; + match result { + ServerResult::CallToolResult(result) => Ok(CallToolResponse::Complete(result)), + ServerResult::InputRequiredResult(result) => { + Ok(CallToolResponse::InputRequired(result)) + } + // SEP-2663 Tasks extension: the server materialized a task. + ServerResult::CreateTaskResult(result) => Ok(CallToolResponse::Task(result)), + _ => Err(ServiceError::UnexpectedResponse), + } + } + + /// SEP-2663 `tasks/get`: poll the current state of a task. + pub async fn get_task(&self, params: GetTaskParams) -> Result { + let result = self + .send_request(ClientRequest::GetTaskRequest(GetTaskRequest::new(params))) + .await?; + match result { + ServerResult::GetTaskResult(result) => Ok(result), + _ => Err(ServiceError::UnexpectedResponse), + } + } + + /// SEP-2663 `tasks/update`: deliver responses to outstanding in-task + /// input requests. The acknowledgement is eventually consistent. + pub async fn update_task(&self, params: UpdateTaskParams) -> Result<(), ServiceError> { + let result = self + .send_request(ClientRequest::UpdateTaskRequest(UpdateTaskRequest::new( + params, + ))) + .await?; + match result { + ServerResult::TaskAckResult(_) | ServerResult::EmptyResult(_) => Ok(()), + _ => Err(ServiceError::UnexpectedResponse), + } + } + + /// SEP-2663 `tasks/cancel`: signal intent to cancel a task. Cancellation + /// is cooperative; the ack does not guarantee the task stops. + pub async fn cancel_task(&self, params: CancelTaskParams) -> Result<(), ServiceError> { + let result = self + .send_request(ClientRequest::CancelTaskRequest(CancelTaskRequest::new( + params, + ))) + .await?; + match result { + ServerResult::TaskAckResult(_) | ServerResult::EmptyResult(_) => Ok(()), + _ => Err(ServiceError::UnexpectedResponse), + } + } + + /// Send one `prompts/get` request and return either a final result or an MRTR + /// `InputRequiredResult` without driving any follow-up rounds. + pub async fn get_prompt_once( + &self, + params: GetPromptRequestParams, + ) -> Result { + let result = self + .send_request(ClientRequest::GetPromptRequest(GetPromptRequest { + method: Default::default(), + params, + extensions: Default::default(), + })) + .await?; + match result { + ServerResult::GetPromptResult(result) => Ok(GetPromptResponse::Complete(result)), + ServerResult::InputRequiredResult(result) => { + Ok(GetPromptResponse::InputRequired(result)) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } + + /// Send one `resources/read` request and return either a final result or an + /// MRTR `InputRequiredResult` without driving any follow-up rounds. + pub async fn read_resource_once( + &self, + params: ReadResourceRequestParams, + ) -> Result { + let cache_key = resource_read_cache_key(¶ms); + if let Some(key) = cache_key.as_deref() + && let Some(ServerResult::ReadResourceResult(result)) = self.cached_response(key).await + { + return Ok(ReadResourceResponse::Complete(result)); + } + + let generation = self.capture_response_cache_generation().await; + let result = self + .send_request(ClientRequest::ReadResourceRequest(ReadResourceRequest { + method: Default::default(), + params, + extensions: Default::default(), + })) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if let Some(key) = cache_key.as_deref() + && let Some(ServerResult::ReadResourceResult(result)) = + self.stale_cached_response(key).await + { + return Ok(ReadResourceResponse::Complete(result)); + } + return Err(error); + } + }; + match result { + ServerResult::ReadResourceResult(result) => { + self.cache_result( + cache_key, + result.ttl_ms, + result.cache_scope, + generation, + ServerResult::ReadResourceResult(result.clone()), + ) + .await; + Ok(ReadResourceResponse::Complete(result)) + } + ServerResult::InputRequiredResult(result) => { + Ok(ReadResourceResponse::InputRequired(result)) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } + method!(peer_req complete CompleteRequest(CompleteRequestParams) => CompleteResult); method!( #[deprecated( @@ -370,14 +1508,232 @@ impl Peer { peer_req set_level SetLevelRequest(SetLevelRequestParams) ); method!(peer_req get_prompt GetPromptRequest(GetPromptRequestParams) => GetPromptResult); - method!(peer_req list_prompts ListPromptsRequest(PaginatedRequestParams)? => ListPromptsResult); - method!(peer_req list_resources ListResourcesRequest(PaginatedRequestParams)? => ListResourcesResult); - method!(peer_req list_resource_templates ListResourceTemplatesRequest(PaginatedRequestParams)? => ListResourceTemplatesResult); - method!(peer_req read_resource ReadResourceRequest(ReadResourceRequestParams) => ReadResourceResult); - method!(peer_req subscribe SubscribeRequest(SubscribeRequestParams) ); - method!(peer_req unsubscribe UnsubscribeRequest(UnsubscribeRequestParams)); + method!( + #[deprecated( + note = "resources/subscribe is legacy-only; use Peer::listen for protocol version 2026-07-28" + )] + peer_req subscribe SubscribeRequest(SubscribeRequestParams) + ); + method!( + #[deprecated( + note = "resources/unsubscribe is legacy-only; cancel the Subscription handle instead" + )] + peer_req unsubscribe UnsubscribeRequest(UnsubscribeRequestParams) + ); method!(peer_req call_tool CallToolRequest(CallToolRequestParams) => CallToolResult); - method!(peer_req list_tools ListToolsRequest(PaginatedRequestParams)? => ListToolsResult); + + pub async fn list_prompts( + &self, + params: Option, + ) -> Result { + let cache_key = list_response_cache_key(PROMPT_LIST_CACHE_PREFIX, ¶ms); + if let Some(ServerResult::ListPromptsResult(result)) = + self.cached_response(&cache_key).await + { + return Ok(result); + } + let generation = self.capture_response_cache_generation().await; + let uses_cursor = request_uses_cursor(¶ms); + let result = self + .send_request(ClientRequest::ListPromptsRequest(ListPromptsRequest { + method: Default::default(), + params, + extensions: Default::default(), + })) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if uses_cursor { + self.invalidate_prompt_cache().await; + return Err(error); + } + if let Some(ServerResult::ListPromptsResult(result)) = + self.stale_cached_response(&cache_key).await + { + return Ok(result); + } + return Err(error); + } + }; + match result { + ServerResult::ListPromptsResult(result) => { + self.cache_result( + Some(cache_key), + result.ttl_ms, + result.cache_scope, + generation, + ServerResult::ListPromptsResult(result.clone()), + ) + .await; + Ok(result) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } + + pub async fn list_resources( + &self, + params: Option, + ) -> Result { + let cache_key = list_response_cache_key(RESOURCE_LIST_CACHE_PREFIX, ¶ms); + if let Some(ServerResult::ListResourcesResult(result)) = + self.cached_response(&cache_key).await + { + return Ok(result); + } + let generation = self.capture_response_cache_generation().await; + let uses_cursor = request_uses_cursor(¶ms); + let result = self + .send_request(ClientRequest::ListResourcesRequest(ListResourcesRequest { + method: Default::default(), + params, + extensions: Default::default(), + })) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if uses_cursor { + self.invalidate_cached_responses(RESOURCE_LIST_CACHE_PREFIX) + .await; + return Err(error); + } + if let Some(ServerResult::ListResourcesResult(result)) = + self.stale_cached_response(&cache_key).await + { + return Ok(result); + } + return Err(error); + } + }; + match result { + ServerResult::ListResourcesResult(result) => { + self.cache_result( + Some(cache_key), + result.ttl_ms, + result.cache_scope, + generation, + ServerResult::ListResourcesResult(result.clone()), + ) + .await; + Ok(result) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } + + pub async fn list_resource_templates( + &self, + params: Option, + ) -> Result { + let cache_key = list_response_cache_key(RESOURCE_TEMPLATE_LIST_CACHE_PREFIX, ¶ms); + if let Some(ServerResult::ListResourceTemplatesResult(result)) = + self.cached_response(&cache_key).await + { + return Ok(result); + } + let generation = self.capture_response_cache_generation().await; + let uses_cursor = request_uses_cursor(¶ms); + let result = self + .send_request(ClientRequest::ListResourceTemplatesRequest( + ListResourceTemplatesRequest { + method: Default::default(), + params, + extensions: Default::default(), + }, + )) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if uses_cursor { + self.invalidate_cached_responses(RESOURCE_TEMPLATE_LIST_CACHE_PREFIX) + .await; + return Err(error); + } + if let Some(ServerResult::ListResourceTemplatesResult(result)) = + self.stale_cached_response(&cache_key).await + { + return Ok(result); + } + return Err(error); + } + }; + match result { + ServerResult::ListResourceTemplatesResult(result) => { + self.cache_result( + Some(cache_key), + result.ttl_ms, + result.cache_scope, + generation, + ServerResult::ListResourceTemplatesResult(result.clone()), + ) + .await; + Ok(result) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } + + pub async fn read_resource( + &self, + params: ReadResourceRequestParams, + ) -> Result { + match self.read_resource_once(params).await? { + ReadResourceResponse::Complete(result) => Ok(result), + ReadResourceResponse::InputRequired(_) => Err(ServiceError::UnexpectedResponse), + } + } + + pub async fn list_tools( + &self, + params: Option, + ) -> Result { + let cache_key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + if let Some(ServerResult::ListToolsResult(result)) = self.cached_response(&cache_key).await + { + return Ok(result); + } + let generation = self.capture_response_cache_generation().await; + let uses_cursor = request_uses_cursor(¶ms); + let result = self + .send_request(ClientRequest::ListToolsRequest(ListToolsRequest { + method: Default::default(), + params, + extensions: Default::default(), + })) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if uses_cursor { + self.invalidate_tool_cache().await; + return Err(error); + } + if let Some(ServerResult::ListToolsResult(result)) = + self.stale_cached_response(&cache_key).await + { + return Ok(result); + } + return Err(error); + } + }; + match result { + ServerResult::ListToolsResult(result) => { + self.cache_result( + Some(cache_key), + result.ttl_ms, + result.cache_scope, + generation, + ServerResult::ListToolsResult(result.clone()), + ) + .await; + Ok(result) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } method!(peer_not notify_cancelled CancelledNotification(CancelledNotificationParam)); method!(peer_not notify_progress ProgressNotification(ProgressNotificationParam)); @@ -468,7 +1824,7 @@ impl Peer { /// /// # Arguments /// * `prompt_name` - Name of the prompt being completed - /// * `argument_name` - Name of the argument being completed + /// * `argument_name` - Name of the argument being completed /// * `current_value` - Current partial value of the argument /// * `context` - Optional context with previously resolved arguments /// @@ -558,3 +1914,813 @@ impl Peer { Ok(completion.values) } } + +impl RunningService +where + S: Service, +{ + /// Send one `tools/call` request without driving MRTR follow-up rounds. + pub async fn call_tool_once( + &self, + params: CallToolRequestParams, + ) -> Result { + self.peer.call_tool_once(params).await + } + + /// Send one `prompts/get` request without driving MRTR follow-up rounds. + pub async fn get_prompt_once( + &self, + params: GetPromptRequestParams, + ) -> Result { + self.peer.get_prompt_once(params).await + } + + /// Send one `resources/read` request without driving MRTR follow-up rounds. + pub async fn read_resource_once( + &self, + params: ReadResourceRequestParams, + ) -> Result { + self.peer.read_resource_once(params).await + } + + /// High-level `tools/call` helper that automatically fulfils SEP-2322 + /// `input_required` rounds through the local [`ClientHandler`](crate::ClientHandler) service. + /// + /// # Errors + /// + /// Returns [`ServiceError::InputRequiredRoundsExceeded`] if the peer does + /// not produce a final [`CallToolResult`] within the default MRTR round cap. + /// Other transport, protocol, and local input-handler errors are propagated. + pub async fn call_tool( + &self, + params: CallToolRequestParams, + ) -> Result { + self.call_tool_with_mrtr_max_rounds(params, DEFAULT_MRTR_MAX_ROUNDS) + .await + } + + /// Same as [`Self::call_tool`], with an explicit MRTR round cap. + /// + /// # Errors + /// + /// Returns [`ServiceError::InputRequiredRoundsExceeded`] once `max_rounds` + /// `input_required` responses have been driven without receiving a final + /// [`CallToolResult`]. Other transport, protocol, and local input-handler + /// errors are propagated. + pub async fn call_tool_with_mrtr_max_rounds( + &self, + mut params: CallToolRequestParams, + max_rounds: usize, + ) -> Result { + let mut state_only_rounds = 0usize; + for _round in 0..max_rounds { + match self.peer.call_tool_once(params.clone()).await? { + CallToolResponse::Complete(result) => return Ok(result), + CallToolResponse::InputRequired(result) => { + let (input_responses, request_state) = self + .prepare_input_required_retry(result, &mut state_only_rounds) + .await?; + params.input_responses = input_responses; + params.request_state = request_state; + } + // SEP-2663: this helper does not drive the task polling + // lifecycle. Callers that declare the tasks extension + // capability should use `call_tool_once` and poll `tasks/get`. + CallToolResponse::Task(_) => return Err(ServiceError::UnexpectedResponse), + } + } + Err(ServiceError::InputRequiredRoundsExceeded { max_rounds }) + } + + /// High-level `prompts/get` helper that automatically fulfils SEP-2322 + /// `input_required` rounds through the local [`ClientHandler`](crate::ClientHandler) service. + /// + /// # Errors + /// + /// Returns [`ServiceError::InputRequiredRoundsExceeded`] if the peer does + /// not produce a final [`GetPromptResult`] within the default MRTR round cap. + /// Other transport, protocol, and local input-handler errors are propagated. + pub async fn get_prompt( + &self, + params: GetPromptRequestParams, + ) -> Result { + self.get_prompt_with_mrtr_max_rounds(params, DEFAULT_MRTR_MAX_ROUNDS) + .await + } + + /// Same as [`Self::get_prompt`], with an explicit MRTR round cap. + /// + /// # Errors + /// + /// Returns [`ServiceError::InputRequiredRoundsExceeded`] once `max_rounds` + /// `input_required` responses have been driven without receiving a final + /// [`GetPromptResult`]. Other transport, protocol, and local input-handler + /// errors are propagated. + pub async fn get_prompt_with_mrtr_max_rounds( + &self, + mut params: GetPromptRequestParams, + max_rounds: usize, + ) -> Result { + let mut state_only_rounds = 0usize; + for _round in 0..max_rounds { + match self.peer.get_prompt_once(params.clone()).await? { + GetPromptResponse::Complete(result) => return Ok(result), + GetPromptResponse::InputRequired(result) => { + let (input_responses, request_state) = self + .prepare_input_required_retry(result, &mut state_only_rounds) + .await?; + params.input_responses = input_responses; + params.request_state = request_state; + } + } + } + Err(ServiceError::InputRequiredRoundsExceeded { max_rounds }) + } + + /// High-level `resources/read` helper that automatically fulfils SEP-2322 + /// `input_required` rounds through the local [`ClientHandler`](crate::ClientHandler) service. + /// + /// # Errors + /// + /// Returns [`ServiceError::InputRequiredRoundsExceeded`] if the peer does + /// not produce a final [`ReadResourceResult`] within the default MRTR round + /// cap. Other transport, protocol, and local input-handler errors are + /// propagated. + pub async fn read_resource( + &self, + params: ReadResourceRequestParams, + ) -> Result { + self.read_resource_with_mrtr_max_rounds(params, DEFAULT_MRTR_MAX_ROUNDS) + .await + } + + /// Same as [`Self::read_resource`], with an explicit MRTR round cap. + /// + /// # Errors + /// + /// Returns [`ServiceError::InputRequiredRoundsExceeded`] once `max_rounds` + /// `input_required` responses have been driven without receiving a final + /// [`ReadResourceResult`]. Other transport, protocol, and local input-handler + /// errors are propagated. + pub async fn read_resource_with_mrtr_max_rounds( + &self, + mut params: ReadResourceRequestParams, + max_rounds: usize, + ) -> Result { + let mut state_only_rounds = 0usize; + for _round in 0..max_rounds { + match self.peer.read_resource_once(params.clone()).await? { + ReadResourceResponse::Complete(result) => return Ok(result), + ReadResourceResponse::InputRequired(result) => { + let (input_responses, request_state) = self + .prepare_input_required_retry(result, &mut state_only_rounds) + .await?; + params.input_responses = input_responses; + params.request_state = request_state; + } + } + } + Err(ServiceError::InputRequiredRoundsExceeded { max_rounds }) + } + + async fn prepare_input_required_retry( + &self, + result: InputRequiredResult, + state_only_rounds: &mut usize, + ) -> Result<(Option, Option), ServiceError> { + let had_input_requests = result + .input_requests + .as_ref() + .is_some_and(|requests| !requests.is_empty()); + if !had_input_requests && result.request_state.is_none() { + return Err(ServiceError::UnexpectedResponse); + } + + let responses = self + .fulfill_input_requests(result.input_requests.unwrap_or_default()) + .await?; + if had_input_requests { + *state_only_rounds = 0; + } else { + Self::sleep_state_only_round(*state_only_rounds).await; + *state_only_rounds += 1; + } + + Ok(( + (!responses.is_empty()).then_some(responses), + result.request_state, + )) + } + + async fn fulfill_input_requests( + &self, + requests: crate::model::InputRequests, + ) -> Result { + let responses = futures::future::try_join_all( + requests + .into_iter() + .map(|(key, request)| self.fulfill_input_request(key, request)), + ) + .await?; + Ok(responses.into_iter().collect()) + } + + async fn fulfill_input_request( + &self, + key: String, + request: InputRequest, + ) -> Result<(String, serde_json::Value), ServiceError> { + let response = match request { + InputRequest::CreateMessage(request) => { + let mut request = ServerRequest::CreateMessageRequest(request); + let context = self.input_request_context(&key, &mut request); + match self + .service + .handle_request(request, context) + .await + .map_err(ServiceError::McpError)? + { + ClientResult::CreateMessageResult(result) => { + serde_json::to_value(result).map_err(Self::serde_to_service_error)? + } + _ => return Err(ServiceError::UnexpectedResponse), + } + } + InputRequest::Elicitation(request) => { + let mut request = ServerRequest::ElicitRequest(request); + let context = self.input_request_context(&key, &mut request); + match self + .service + .handle_request(request, context) + .await + .map_err(ServiceError::McpError)? + { + ClientResult::ElicitResult(result) => { + serde_json::to_value(result).map_err(Self::serde_to_service_error)? + } + _ => return Err(ServiceError::UnexpectedResponse), + } + } + InputRequest::ListRoots(request) => { + let mut request = ServerRequest::ListRootsRequest(request); + let context = self.input_request_context(&key, &mut request); + match self + .service + .handle_request(request, context) + .await + .map_err(ServiceError::McpError)? + { + ClientResult::ListRootsResult(result) => { + serde_json::to_value(result).map_err(Self::serde_to_service_error)? + } + _ => return Err(ServiceError::UnexpectedResponse), + } + } + }; + Ok((key, response)) + } + + fn input_request_context(&self, key: &str, request: &mut T) -> RequestContext + where + T: GetMeta + GetExtensions, + { + let mut meta = Default::default(); + let mut extensions = Default::default(); + std::mem::swap(&mut meta, request.get_meta_mut()); + std::mem::swap(&mut extensions, request.extensions_mut()); + RequestContext { + ct: tokio_util::sync::CancellationToken::new(), + id: NumberOrString::String(Arc::from(key)), + peer: self.peer.clone(), + meta, + extensions, + } + } + + async fn sleep_state_only_round(state_only_rounds: usize) { + let millis = (50u64.saturating_mul(1_u64 << state_only_rounds.min(3))).min(250); + tokio::time::sleep(Duration::from_millis(millis)).await; + } + + fn serde_to_service_error(error: serde_json::Error) -> ServiceError { + ServiceError::McpError(ErrorData::internal_error( + format!("failed to serialize MRTR input response: {error}"), + None, + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn server_cancellation_retires_subscription_responder() { + tokio::task::LocalSet::new() + .run_until(async { + for (discover, version) in [ + (false, ProtocolVersion::V_2026_07_28), + (true, ProtocolVersion::V_2026_07_28), + (true, ProtocolVersion::V_2025_11_25), + ] { + tokio::time::timeout( + Duration::from_secs(5), + check_subscription_cancellation(discover, version), + ) + .await + .expect("subscription cancellation timed out"); + } + }) + .await; + } + + async fn check_subscription_cancellation(discover: bool, version: ProtocolVersion) { + use crate::model::{ + GetMeta, PingRequest, ServerInfo, SubscriptionsAcknowledgedNotification, + SubscriptionsAcknowledgedNotificationParams, + }; + + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = + crate::transport::IntoTransport::::into_transport(server_transport); + let info = ServerInfo { + protocol_version: version.clone(), + ..Default::default() + }; + let client = if discover { + let (client, ()) = tokio::join!( + serve_client_with_lifecycle( + (), + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![version.clone()] + } + ), + async { + let Some(ClientJsonRpcMessage::Request(request)) = server.receive().await + else { + panic!("expected discover request"); + }; + assert!(matches!(request.request, ClientRequest::DiscoverRequest(_))); + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult(DiscoverResult::new( + vec![version], + info.capabilities, + )), + request.id, + )) + .await + .unwrap(); + } + ); + client.unwrap() + } else { + super::super::serve_directly::( + (), + client_transport, + Some(info.into()), + ) + }; + let filter = SubscriptionFilter::default(); + let (subscription, ()) = tokio::join!(client.listen(filter.clone()), async { + let Some(ClientJsonRpcMessage::Request(request)) = server.receive().await else { + panic!("expected listen request"); + }; + assert!(matches!( + request.request, + ClientRequest::SubscriptionsListenRequest(_) + )); + let mut ack: ServerNotification = SubscriptionsAcknowledgedNotification::new( + SubscriptionsAcknowledgedNotificationParams::new(filter), + ) + .into(); + ack.get_meta_mut().set_subscription_id(request.id); + server + .send(ServerJsonRpcMessage::notification(ack)) + .await + .unwrap(); + }); + let mut subscription = subscription.unwrap(); + let id = subscription.id().clone(); + let mut ordinary = client + .send_cancellable_request( + PingRequest { + method: Default::default(), + extensions: Default::default(), + } + .into(), + PeerRequestOptions::no_options(), + ) + .await + .unwrap(); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Request(_)) + )); + for unknown in [ + ordinary.id.clone(), + RequestId::String(id.to_string().into()), + RequestId::String("unknown".into()), + ] { + server + .send(ServerJsonRpcMessage::notification( + CancelledNotification::new(CancelledNotificationParam::new( + Some(unknown), + None, + )) + .into(), + )) + .await + .unwrap(); + } + server + .send(ServerJsonRpcMessage::request( + PingRequest { + method: Default::default(), + extensions: Default::default(), + } + .into(), + RequestId::String("before-cancellation".into()), + )) + .await + .unwrap(); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Response(_)) + )); + assert!(matches!( + subscription.request.as_mut().unwrap().rx.try_recv(), + Err(tokio::sync::oneshot::error::TryRecvError::Empty) + )); + assert!(matches!( + ordinary.rx.try_recv(), + Err(tokio::sync::oneshot::error::TryRecvError::Empty) + )); + assert!(matches!( + subscription.notifications.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + server + .send(ServerJsonRpcMessage::notification( + CancelledNotification::new(CancelledNotificationParam::new( + Some(id.clone()), + Some("subscription ended".to_owned()), + )) + .into(), + )) + .await + .unwrap(); + // Ordered input makes the ping response a barrier for cancellation handling. + server + .send(ServerJsonRpcMessage::request( + PingRequest { + method: Default::default(), + extensions: Default::default(), + } + .into(), + RequestId::String("after-cancellation".into()), + )) + .await + .unwrap(); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Response(_)) + )); + assert!( + matches!( + subscription.request.as_mut().unwrap().rx.try_recv(), + Ok(Err(ServiceError::Cancelled { .. })) + ), + "the outbound responder must retire without a final response or disconnect" + ); + assert!(subscription.next().await.unwrap().is_none()); + assert!(matches!( + subscription.end(), + Some(SubscriptionEnd::Cancelled) + )); + assert!(client.peer().subscription_sender(&id).is_none()); + server + .send(ServerJsonRpcMessage::response( + ServerResult::empty(()), + ordinary.id.clone(), + )) + .await + .unwrap(); + assert!(ordinary.await_response().await.is_ok()); + + // Raw listen requests also retire, even without a registered notification channel. + let raw = client + .send_cancellable_request( + ClientRequest::SubscriptionsListenRequest(SubscriptionsListenRequest::new( + SubscriptionsListenRequestParams::new(SubscriptionFilter::default()), + )), + PeerRequestOptions::no_options(), + ) + .await + .unwrap(); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Request(_)) + )); + server + .send(ServerJsonRpcMessage::notification( + CancelledNotification::new(CancelledNotificationParam::new( + Some(raw.id.clone()), + None, + )) + .into(), + )) + .await + .unwrap(); + assert!(matches!( + raw.await_response().await, + Err(ServiceError::Cancelled { .. }) + )); + client.cancel().await.unwrap(); + } + + #[tokio::test] + async fn auto_startup_falls_back_when_discover_is_ignored() { + use crate::model::{InitializeResult, ServerCapabilities}; + + tokio::task::LocalSet::new() + .run_until(async { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = + crate::transport::IntoTransport::::into_transport( + server_transport, + ); + let server_task = tokio::task::spawn_local(async move { + let ClientJsonRpcMessage::Request(discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected discover request"); + }; + assert!(matches!( + discover.request, + ClientRequest::DiscoverRequest(_) + )); + + let ClientJsonRpcMessage::Request(initialize) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected initialize request"); + }; + assert!(matches!( + initialize.request, + ClientRequest::InitializeRequest(_) + )); + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult(InitializeResult::new( + ServerCapabilities::default(), + )), + initialize.id, + )) + .await + .expect("send initialize response"); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + }); + + let client_transport = + crate::transport::IntoTransport::::into_transport( + client_transport, + ); + let client = serve_client_with_ct_inner( + (), + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + CancellationToken::new(), + Duration::from_millis(25), + ) + .await + .expect("auto client should fall back after discover timeout"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); + }) + .await; + } + + fn disconnected_peer() -> Peer { + let (peer, receiver) = + Peer::::new(Arc::new(AtomicU32RequestIdProvider::default()), None); + drop(receiver); + peer + } + + fn tools_result(ttl_ms: Option, cache_scope: Option) -> ListToolsResult { + let mut result = ListToolsResult::with_all_items(Vec::new()); + result.ttl_ms = ttl_ms; + result.cache_scope = cache_scope; + result + } + + #[tokio::test] + async fn fresh_cached_page_is_served_without_transport_io() { + let peer = disconnected_peer(); + let params = None::; + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + let expected = tools_result(Some(5_000), Some(CacheScope::Public)); + peer.cache_response( + key, + ServerResult::ListToolsResult(expected.clone()), + expected.ttl_ms, + expected.cache_scope, + ) + .await; + + assert_eq!(peer.list_tools(params).await.unwrap(), expected); + } + + #[tokio::test] + async fn expired_entry_falls_through_to_the_transport() { + let peer = disconnected_peer(); + peer.set_response_cache_config( + ClientCacheConfig::default().with_serve_stale_on_error(false), + ) + .await; + let params = None::; + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + peer.cache_response( + key, + ServerResult::ListToolsResult(tools_result(Some(1), Some(CacheScope::Public))), + Some(1), + Some(CacheScope::Public), + ) + .await; + tokio::time::sleep(Duration::from_millis(5)).await; + + assert!(matches!( + peer.list_tools(params).await, + Err(ServiceError::TransportClosed) + )); + } + + #[tokio::test] + async fn private_entries_are_isolated_between_authorization_partitions() { + let peer = disconnected_peer(); + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, &None); + + peer.set_response_cache_config( + ClientCacheConfig::default().with_private_partition("auth-a"), + ) + .await; + peer.cache_response( + key.clone(), + ServerResult::ListToolsResult(tools_result(Some(5_000), Some(CacheScope::Private))), + Some(5_000), + Some(CacheScope::Private), + ) + .await; + assert!(peer.cached_response(&key).await.is_some()); + + // Switching to a different authorization context must not expose the + // first partition's private entry. + peer.set_response_cache_config( + ClientCacheConfig::default().with_private_partition("auth-b"), + ) + .await; + assert!(peer.cached_response(&key).await.is_none()); + } + + #[tokio::test] + async fn list_change_notification_discards_every_cached_page() { + let peer = disconnected_peer(); + for cursor in [None, Some("page-a".into()), Some("page-b".into())] { + let params = + cursor.map(|cursor| PaginatedRequestParams::default().with_cursor(Some(cursor))); + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + peer.cache_response( + key, + ServerResult::ListToolsResult(tools_result(Some(5_000), Some(CacheScope::Public))), + Some(5_000), + Some(CacheScope::Public), + ) + .await; + } + + peer.invalidate_tool_cache().await; + + for cursor in [None, Some("page-a".into()), Some("page-b".into())] { + let params = + cursor.map(|cursor| PaginatedRequestParams::default().with_cursor(Some(cursor))); + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + assert!(peer.cached_response(&key).await.is_none()); + } + } + + #[tokio::test] + async fn expired_entry_is_served_when_refetch_fails() { + let peer = disconnected_peer(); + let params = None::; + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + let expected = tools_result(Some(1), Some(CacheScope::Public)); + peer.cache_response( + key, + ServerResult::ListToolsResult(expected.clone()), + Some(1), + Some(CacheScope::Public), + ) + .await; + tokio::time::sleep(Duration::from_millis(5)).await; + + assert_eq!(peer.list_tools(params).await.unwrap(), expected); + } + + #[tokio::test] + async fn discover_serves_a_fresh_cached_response_without_transport_io() { + let peer = disconnected_peer(); + let meta = RequestMetaObject::default(); + let key = discover_cache_key(); + let expected = DiscoverResult::new(vec![ProtocolVersion::default()], Default::default()) + .with_server_info(crate::model::Implementation::from_build_env()) + .with_ttl_ms(5_000) + .with_cache_scope(CacheScope::Public); + peer.cache_response( + key, + ServerResult::DiscoverResult(expected.clone()), + Some(5_000), + Some(CacheScope::Public), + ) + .await; + + assert_eq!(peer.discover(meta).await.unwrap(), expected); + } +} + +#[cfg(test)] +mod sep2260_association_tests { + use super::*; + use crate::{ + model::{ + CreateMessageRequest, CreateMessageRequestParams, SamplingMessage, ServerCapabilities, + }, + service::PeerRequestAssociation, + }; + + fn sampling_request() -> ServerRequest { + ServerRequest::CreateMessageRequest(CreateMessageRequest::new( + CreateMessageRequestParams::new(vec![SamplingMessage::user_text("hi")], 16), + )) + } + + fn server_info(version: ProtocolVersion) -> ServerPeerInfo { + ServerPeerInfo::new(version, ServerCapabilities::default()) + } + + fn enforce( + info: &ServerPeerInfo, + association: PeerRequestAssociation, + ) -> Result<(), ErrorData> { + RoleClient::enforce_peer_request_association(&sampling_request(), Some(info), association) + } + + #[test] + fn strict_rejects_unassociated() { + let info = server_info(ProtocolVersion::V_2026_07_28); + let err = enforce(&info, PeerRequestAssociation::Unassociated).unwrap_err(); + assert_eq!(err.code, crate::model::ErrorCode::INVALID_PARAMS); + } + + #[test] + fn strict_accepts_associated() { + let info = server_info(ProtocolVersion::V_2026_07_28); + assert!(enforce(&info, PeerRequestAssociation::Associated).is_ok()); + } + + #[test] + fn strict_unknown_falls_back_to_coarse_check() { + let info = server_info(ProtocolVersion::V_2026_07_28); + assert!( + enforce( + &info, + PeerRequestAssociation::Unknown { + has_pending_outbound_request: true + } + ) + .is_ok() + ); + assert!( + enforce( + &info, + PeerRequestAssociation::Unknown { + has_pending_outbound_request: false + } + ) + .is_err() + ); + } + + #[test] + fn legacy_protocol_accepts_even_unassociated() { + let info = server_info(ProtocolVersion::V_2025_11_25); + assert!(enforce(&info, PeerRequestAssociation::Unassociated).is_ok()); + } +} diff --git a/crates/rmcp/src/service/client/cache.rs b/crates/rmcp/src/service/client/cache.rs new file mode 100644 index 0000000..f7137cc --- /dev/null +++ b/crates/rmcp/src/service/client/cache.rs @@ -0,0 +1,401 @@ +use std::{ + collections::HashMap, + sync::Arc, + time::{Duration, Instant}, +}; + +use super::RoleClient; +use crate::{ + model::CacheScope, + service::{Peer, ServiceRole}, +}; + +/// Maximum server-provided cache TTL honoured by the client response cache. +pub const MAX_CLIENT_CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60); + +/// Configuration for the built-in MCP client response cache. +/// +/// A cache is allocated per client [`Peer`]. Public responses may be reused +/// throughout that client connection. Private responses are additionally +/// partitioned by `private_partition`; changing the partition drops every +/// private entry while preserving public entries. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub struct ClientCacheConfig { + /// Enables cache reads and writes. + pub enabled: bool, + /// TTL used when a backwards-compatible server omits `ttlMs`. + /// + /// The default is zero, which leaves such responses immediately stale. + pub default_ttl: Duration, + /// Upper bound applied to both server-provided and default TTLs. + pub max_ttl: Duration, + /// Stable opaque identity for the current authorization context. + /// + /// A single-principal client may leave this unset because each client owns + /// its own in-memory store. Gateways or clients that change principals on an + /// existing connection should set this value and update it whenever the + /// authorization context changes. + pub private_partition: Option, + /// Maximum number of responses retained by the in-memory cache. + /// + /// A value of zero disables the size limit. + pub max_entries: usize, + /// Serves an expired cached response when a re-fetch fails. + /// + /// SEP-2549 permits clients to serve stale responses if errors occur while + /// re-fetching (for example, network issues or server downtime). When this + /// is enabled the client retains expired entries so it can fall back to the + /// last known response instead of surfacing the transport or server error. + /// A successful re-fetch always overwrites the stale entry. + pub serve_stale_on_error: bool, +} + +impl Default for ClientCacheConfig { + fn default() -> Self { + Self { + enabled: true, + default_ttl: Duration::ZERO, + max_ttl: MAX_CLIENT_CACHE_TTL, + private_partition: None, + max_entries: 512, + serve_stale_on_error: true, + } + } +} + +impl ClientCacheConfig { + /// Returns a configuration that disables all cache reads and writes. + pub fn disabled() -> Self { + Self { + enabled: false, + ..Self::default() + } + } + + /// Sets the TTL used when a response omits `ttlMs`. + pub fn with_default_ttl(mut self, default_ttl: Duration) -> Self { + self.default_ttl = default_ttl; + self + } + + /// Sets the maximum TTL the client will honour. + pub fn with_max_ttl(mut self, max_ttl: Duration) -> Self { + self.max_ttl = max_ttl; + self + } + + /// Sets the stable partition for private responses. + pub fn with_private_partition(mut self, partition: impl Into) -> Self { + self.private_partition = Some(partition.into()); + self + } + + /// Sets the maximum number of retained responses. + pub fn with_max_entries(mut self, max_entries: usize) -> Self { + self.max_entries = max_entries; + self + } + + /// Controls whether an expired response may be served when a re-fetch fails. + pub fn with_serve_stale_on_error(mut self, serve_stale_on_error: bool) -> Self { + self.serve_stale_on_error = serve_stale_on_error; + self + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +enum CachePartition { + Public, + Private(Arc), +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct CacheKey { + logical_key: String, + partition: CachePartition, +} + +#[derive(Debug, Clone)] +struct CachedPeerResponse { + value: T, + expires_at: Instant, + inserted_at: Instant, + scope: CacheScope, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct CacheGeneration(u64); + +#[derive(Debug)] +pub(crate) struct PeerResponseCacheState { + entries: HashMap>, + config: ClientCacheConfig, + generation: u64, +} + +impl Default for PeerResponseCacheState { + fn default() -> Self { + Self { + entries: HashMap::new(), + config: ClientCacheConfig::default(), + generation: 0, + } + } +} + +impl PeerResponseCacheState { + fn trim_to_limit(&mut self) { + while self.config.max_entries > 0 && self.entries.len() > self.config.max_entries { + let Some(oldest_key) = self + .entries + .iter() + .min_by_key(|(_, entry)| entry.inserted_at) + .map(|(key, _)| key.clone()) + else { + break; + }; + self.entries.remove(&oldest_key); + } + } +} + +pub(crate) type PeerResponseCache = Arc>>; + +impl Peer { + fn private_partition(config: &ClientCacheConfig) -> Arc { + Arc::from(config.private_partition.as_deref().unwrap_or("connection")) + } + + fn cache_key(logical_key: &str, partition: CachePartition) -> CacheKey { + CacheKey { + logical_key: logical_key.to_owned(), + partition, + } + } + + fn scoped_cache_key( + logical_key: &str, + scope: CacheScope, + config: &ClientCacheConfig, + ) -> CacheKey { + let partition = match scope { + CacheScope::Public => CachePartition::Public, + CacheScope::Private => CachePartition::Private(Self::private_partition(config)), + }; + Self::cache_key(logical_key, partition) + } + + /// Captures the cache generation before a request crosses the transport. + /// + /// Any configuration change, explicit clear, or notification invalidation + /// advances the generation. A response from an older generation is not + /// written back, preventing an in-flight stale response from undoing an + /// invalidation. + pub(crate) async fn capture_response_cache_generation(&self) -> CacheGeneration { + CacheGeneration(self.response_cache.read().await.generation) + } + + /// Returns a fresh cached response, preferring the current private partition + /// before the public partition. + /// + /// Expired entries are removed on access unless `serve_stale_on_error` is + /// enabled, in which case they are retained so a later re-fetch failure can + /// fall back to them via [`Peer::stale_cached_response`]. + pub(crate) async fn cached_response(&self, logical_key: &str) -> Option { + let now = Instant::now(); + let mut cache = self.response_cache.write().await; + if !cache.config.enabled { + return None; + } + let keep_stale = cache.config.serve_stale_on_error; + + let private_key = Self::cache_key( + logical_key, + CachePartition::Private(Self::private_partition(&cache.config)), + ); + let private_fresh = cache.entries.get(&private_key).and_then(|entry| { + (entry.expires_at > now && entry.scope == CacheScope::Private) + .then(|| entry.value.clone()) + }); + if let Some(value) = private_fresh { + return Some(value); + } + if !keep_stale { + cache.entries.remove(&private_key); + } + + let public_key = Self::cache_key(logical_key, CachePartition::Public); + let public_fresh = cache.entries.get(&public_key).and_then(|entry| { + (entry.expires_at > now && entry.scope == CacheScope::Public) + .then(|| entry.value.clone()) + }); + if let Some(value) = public_fresh { + return Some(value); + } + if !keep_stale { + cache.entries.remove(&public_key); + } + None + } + + /// Returns a cached response ignoring its TTL, for use as a fallback when a + /// re-fetch fails (SEP-2549 permits serving stale responses on error). + /// + /// Returns `None` when the cache is disabled or `serve_stale_on_error` is + /// turned off. The private partition is preferred over the public one. The + /// entry is left in place so repeated failures keep serving it until a + /// successful re-fetch overwrites it or a notification invalidates it. + pub(crate) async fn stale_cached_response(&self, logical_key: &str) -> Option { + let cache = self.response_cache.read().await; + if !cache.config.enabled || !cache.config.serve_stale_on_error { + return None; + } + + let private_key = Self::cache_key( + logical_key, + CachePartition::Private(Self::private_partition(&cache.config)), + ); + if let Some(entry) = cache.entries.get(&private_key) + && entry.scope == CacheScope::Private + { + return Some(entry.value.clone()); + } + + let public_key = Self::cache_key(logical_key, CachePartition::Public); + if let Some(entry) = cache.entries.get(&public_key) + && entry.scope == CacheScope::Public + { + return Some(entry.value.clone()); + } + None + } + + /// Stores a response when the configured effective TTL is positive. + /// + /// Missing `cacheScope` is treated as private. This is deliberately more + /// conservative than the model's backwards-compatible wire default and + /// prevents an older or malformed server response from becoming shareable. + pub(crate) async fn cache_response_with_generation( + &self, + logical_key: String, + value: R::PeerResp, + ttl_ms: Option, + cache_scope: Option, + generation: CacheGeneration, + ) { + let now = Instant::now(); + let mut cache = self.response_cache.write().await; + if !cache.config.enabled || generation.0 != cache.generation { + return; + } + + let requested_ttl = ttl_ms + .map(Duration::from_millis) + .unwrap_or(cache.config.default_ttl); + let ttl = requested_ttl.min(cache.config.max_ttl); + if ttl.is_zero() { + return; + } + let Some(expires_at) = now.checked_add(ttl) else { + return; + }; + let scope = cache_scope.unwrap_or(CacheScope::Private); + let target_key = Self::scoped_cache_key(&logical_key, scope, &cache.config); + let opposite_key = match scope { + CacheScope::Public => Self::cache_key( + &logical_key, + CachePartition::Private(Self::private_partition(&cache.config)), + ), + CacheScope::Private => Self::cache_key(&logical_key, CachePartition::Public), + }; + + if !cache.config.serve_stale_on_error { + cache.entries.retain(|_, entry| entry.expires_at > now); + } + cache.entries.remove(&opposite_key); + + if cache.config.max_entries > 0 + && !cache.entries.contains_key(&target_key) + && cache.entries.len() >= cache.config.max_entries + && let Some(oldest_key) = cache + .entries + .iter() + .min_by_key(|(_, entry)| entry.inserted_at) + .map(|(key, _)| key.clone()) + { + cache.entries.remove(&oldest_key); + } + + cache.entries.insert( + target_key, + CachedPeerResponse { + value, + expires_at, + inserted_at: now, + scope, + }, + ); + } + + #[cfg(test)] + pub(crate) async fn cache_response( + &self, + logical_key: String, + value: R::PeerResp, + ttl_ms: Option, + cache_scope: Option, + ) { + let generation = self.capture_response_cache_generation().await; + self.cache_response_with_generation(logical_key, value, ttl_ms, cache_scope, generation) + .await; + } + + pub(crate) async fn invalidate_cached_responses(&self, prefix: &str) { + let mut cache = self.response_cache.write().await; + cache.generation = cache.generation.wrapping_add(1); + cache + .entries + .retain(|key, _| !key.logical_key.starts_with(prefix)); + } +} + +impl Peer { + /// Replaces the response-cache configuration. + /// + /// Changing the private partition invalidates private entries from the old + /// authorization context. Disabling the cache clears every entry. Any + /// configuration change also suppresses writes from requests that were + /// already in flight under the previous configuration. + pub async fn set_response_cache_config(&self, config: ClientCacheConfig) { + let mut cache = self.response_cache.write().await; + let config_changed = cache.config != config; + let partition_changed = cache.config.private_partition != config.private_partition; + let ttl_policy_changed = cache.config.default_ttl != config.default_ttl + || cache.config.max_ttl != config.max_ttl; + cache.config = config; + if config_changed { + cache.generation = cache.generation.wrapping_add(1); + } + if !cache.config.enabled || ttl_policy_changed { + cache.entries.clear(); + } else if partition_changed { + cache + .entries + .retain(|_, entry| entry.scope == CacheScope::Public); + } + cache.trim_to_limit(); + } + + /// Returns a snapshot of the active response-cache configuration. + pub async fn response_cache_config(&self) -> ClientCacheConfig { + self.response_cache.read().await.config.clone() + } + + /// Clears every cached client response without changing the configuration. + pub async fn clear_response_cache(&self) { + let mut cache = self.response_cache.write().await; + cache.generation = cache.generation.wrapping_add(1); + cache.entries.clear(); + } +} diff --git a/crates/rmcp/src/service/server.rs b/crates/rmcp/src/service/server.rs index 4f479b9..3ff229f 100644 --- a/crates/rmcp/src/service/server.rs +++ b/crates/rmcp/src/service/server.rs @@ -1,8 +1,8 @@ // Sampling/Roots/Logging are SEP-2577-deprecated; internal references are expected. #![expect(deprecated)] -use std::borrow::Cow; #[cfg(feature = "elicitation")] use std::collections::HashSet; +use std::{borrow::Cow, sync::Arc}; use thiserror::Error; #[cfg(feature = "elicitation")] @@ -10,20 +10,19 @@ use url::Url; use super::*; #[cfg(feature = "elicitation")] -use crate::model::{ - ElicitRequest, ElicitRequestParams, ElicitResult, ElicitationAction, - ElicitationCompleteNotification, ElicitationResponseNotificationParam, -}; +use crate::model::{ElicitRequest, ElicitRequestParams, ElicitResult, ElicitationAction}; use crate::{ model::{ - CancelledNotification, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage, + CancelledNotification, CancelledNotificationParam, ClientJsonRpcMessage, ClientNotification, ClientRequest, ClientResult, CreateMessageRequest, - CreateMessageRequestParams, CreateMessageResult, EmptyResult, ErrorData, ListRootsRequest, - ListRootsResult, LoggingMessageNotification, LoggingMessageNotificationParam, - ProgressNotification, ProgressNotificationParam, PromptListChangedNotification, - ProtocolVersion, ResourceListChangedNotification, ResourceUpdatedNotification, - ResourceUpdatedNotificationParam, ServerInfo, ServerNotification, ServerRequest, - ServerResult, ToolListChangedNotification, + CreateMessageRequestParams, CreateMessageResult, EmptyResult, ErrorData, + InitializeRequestParams, ListRootsRequest, ListRootsResult, LoggingMessageNotification, + LoggingMessageNotificationParam, ProgressNotification, ProgressNotificationParam, + PromptListChangedNotification, ProtocolVersion, ResourceListChangedNotification, + ResourceUpdatedNotification, ResourceUpdatedNotificationParam, ServerConfig, + ServerNotification, ServerRequest, ServerResult, SubscriptionFilter, + SubscriptionsAcknowledgedNotification, SubscriptionsAcknowledgedNotificationParams, + ToolListChangedNotification, }, transport::DynamicTransportError, }; @@ -39,11 +38,43 @@ impl ServiceRole for RoleServer { type PeerReq = ClientRequest; type PeerResp = ClientResult; type PeerNot = ClientNotification; - type Info = ServerInfo; - type PeerInfo = ClientInfo; + type Info = ServerConfig; + type PeerInfo = InitializeRequestParams; type InitializeError = ServerInitializeError; const IS_CLIENT: bool = false; + + fn peer_cancelled_params(notification: &Self::PeerNot) -> Option<&CancelledNotificationParam> { + match notification { + ClientNotification::CancelledNotification(notification) => Some(¬ification.params), + _ => None, + } + } + + fn enforce_request_association( + request: &Self::Req, + peer_info: Option<&Self::PeerInfo>, + in_request_handler_scope: bool, + ) -> Result<(), ServiceError> { + let restricted = matches!( + request, + ServerRequest::CreateMessageRequest(_) + | ServerRequest::ListRootsRequest(_) + | ServerRequest::ElicitRequest(_) + ); + if !restricted { + return Ok(()); + } + let strict = + peer_info.is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28); + if strict && !in_request_handler_scope { + return Err(ServiceError::McpError(ErrorData::invalid_request( + "SEP-2260: server-to-client requests must be associated with an originating client request", + None, + ))); + } + Ok(()) + } } /// It represents the error that may occur when serving the server. @@ -55,13 +86,6 @@ pub enum ServerInitializeError { #[error("expect initialized request, but received: {0:?}")] ExpectedInitializeRequest(Option), - #[deprecated( - since = "1.4.0", - note = "The server no longer gates on the initialized notification. This variant is never constructed and will be removed in a future major release." - )] - #[error("expect initialized notification, but received: {0:?}")] - ExpectedInitializedNotification(Option), - #[error("connection closed: {0}")] ConnectionClosed(String), @@ -71,13 +95,6 @@ pub enum ServerInitializeError { #[error("initialize failed: {0}")] InitializeFailed(ErrorData), - #[deprecated( - since = "1.8.0", - note = "Negotiation now falls back to the server-configured version. This variant is never constructed and will be removed in a future major release." - )] - #[error("unsupported protocol version: {0}")] - UnsupportedProtocolVersion(ProtocolVersion), - #[error("Send message error {error}, when {context}")] TransportError { error: DynamicTransportError, @@ -101,6 +118,289 @@ impl ServerInitializeError { } pub type ClientSink = Peer; +/// Failure to send a notification through a [`SubscriptionSink`]. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum SubscriptionSendError { + #[error("subscription is no longer active")] + SubscriptionClosed, + #[error("notification is not allowed on a subscription stream: {0}")] + UnsupportedNotification(&'static str), + #[error("notification was not accepted for this subscription: {0}")] + NotificationNotAccepted(&'static str), + #[error(transparent)] + Service(#[from] ServiceError), +} + +/// A server-side notification sink scoped to one `subscriptions/listen` request. +/// +/// The sink applies the accepted filter and adds the subscription request ID to +/// every notification it sends. +#[derive(Debug, Clone)] +pub struct SubscriptionSink { + peer: Peer, + id: RequestId, + accepted: Arc, + active: CancellationToken, +} + +impl SubscriptionSink { + fn new( + peer: Peer, + id: RequestId, + accepted: Arc, + active: CancellationToken, + ) -> Self { + Self { + peer, + id, + accepted, + active, + } + } + + /// Return the JSON-RPC ID of the originating listen request. + pub fn id(&self) -> &RequestId { + &self.id + } + + /// Return the filter accepted for this subscription. + pub fn accepted(&self) -> &SubscriptionFilter { + self.accepted.as_ref() + } + + /// Send an allowed change notification with subscription metadata attached. + /// + /// # Errors + /// + /// Returns [`SubscriptionSendError::SubscriptionClosed`] after the request + /// ends, a filter error for disallowed notifications, or a transport error. + pub async fn send( + &self, + mut notification: ServerNotification, + ) -> Result<(), SubscriptionSendError> { + if self.active.is_cancelled() { + return Err(SubscriptionSendError::SubscriptionClosed); + } + match ¬ification { + ServerNotification::ToolListChangedNotification(_) => { + if self.accepted.tools_list_changed != Some(true) { + return Err(SubscriptionSendError::NotificationNotAccepted( + "notifications/tools/list_changed", + )); + } + } + ServerNotification::PromptListChangedNotification(_) => { + if self.accepted.prompts_list_changed != Some(true) { + return Err(SubscriptionSendError::NotificationNotAccepted( + "notifications/prompts/list_changed", + )); + } + } + ServerNotification::ResourceListChangedNotification(_) => { + if self.accepted.resources_list_changed != Some(true) { + return Err(SubscriptionSendError::NotificationNotAccepted( + "notifications/resources/list_changed", + )); + } + } + ServerNotification::ResourceUpdatedNotification(update) => { + let accepted = self + .accepted + .resource_subscriptions + .as_ref() + .is_some_and(|uris| uris.contains(&update.params.uri)); + if !accepted { + return Err(SubscriptionSendError::NotificationNotAccepted( + "notifications/resources/updated", + )); + } + } + ServerNotification::SubscriptionsAcknowledgedNotification(_) => { + return Err(SubscriptionSendError::UnsupportedNotification( + "notifications/subscriptions/acknowledged", + )); + } + ServerNotification::CancelledNotification(_) => { + return Err(SubscriptionSendError::UnsupportedNotification( + "notifications/cancelled", + )); + } + ServerNotification::ProgressNotification(_) => { + return Err(SubscriptionSendError::UnsupportedNotification( + "notifications/progress", + )); + } + ServerNotification::LoggingMessageNotification(_) => { + return Err(SubscriptionSendError::UnsupportedNotification( + "notifications/message", + )); + } + // SEP-2663 task status notifications are not yet routable through + // `subscriptions/listen`: `SubscriptionFilter` has no `taskIds` + // field yet (the upstream conformance check for this flow is also + // still skipped, pending the subscriptions/listen rewrite). + // Clients currently observe task state by polling `tasks/get`. + ServerNotification::TaskStatusNotification(_) => { + return Err(SubscriptionSendError::UnsupportedNotification( + "notifications/tasks", + )); + } + ServerNotification::CustomNotification(_) => { + return Err(SubscriptionSendError::UnsupportedNotification( + "custom notification", + )); + } + } + + notification + .get_meta_mut() + .set_subscription_id(self.id.clone()); + self.peer.send_notification(notification).await?; + Ok(()) + } + + /// Send `notifications/tools/list_changed`. + /// + /// # Errors + /// + /// See [`Self::send`]. + pub async fn notify_tool_list_changed(&self) -> Result<(), SubscriptionSendError> { + self.send(ServerNotification::ToolListChangedNotification( + ToolListChangedNotification { + method: Default::default(), + extensions: Default::default(), + }, + )) + .await + } + + /// Send `notifications/prompts/list_changed`. + /// + /// # Errors + /// + /// See [`Self::send`]. + pub async fn notify_prompt_list_changed(&self) -> Result<(), SubscriptionSendError> { + self.send(ServerNotification::PromptListChangedNotification( + PromptListChangedNotification { + method: Default::default(), + extensions: Default::default(), + }, + )) + .await + } + + /// Send `notifications/resources/list_changed`. + /// + /// # Errors + /// + /// See [`Self::send`]. + pub async fn notify_resource_list_changed(&self) -> Result<(), SubscriptionSendError> { + self.send(ServerNotification::ResourceListChangedNotification( + ResourceListChangedNotification { + method: Default::default(), + extensions: Default::default(), + }, + )) + .await + } + + /// Send `notifications/resources/updated` for an accepted URI. + /// + /// # Errors + /// + /// See [`Self::send`]. + pub async fn notify_resource_updated( + &self, + uri: impl Into, + ) -> Result<(), SubscriptionSendError> { + self.send(ServerNotification::ResourceUpdatedNotification( + ResourceUpdatedNotification::new(ResourceUpdatedNotificationParam::new(uri)), + )) + .await + } +} + +/// Context for one established server-side notification subscription. +/// +/// The acknowledgment has already been sent before this context is handed to +/// [`ServerHandler::listen`](crate::ServerHandler::listen). +#[derive(Debug)] +pub struct SubscriptionContext { + request: RequestContext, + requested: SubscriptionFilter, + accepted: Arc, + sink: SubscriptionSink, + _active_guard: DropGuard, +} + +impl SubscriptionContext { + pub(crate) async fn establish( + request: RequestContext, + requested: SubscriptionFilter, + accepted: SubscriptionFilter, + ) -> Result { + let active = request.ct.child_token(); + let accepted = Arc::new(accepted); + let sink = SubscriptionSink::new( + request.peer.clone(), + request.id.clone(), + accepted.clone(), + active.clone(), + ); + let mut acknowledgment = SubscriptionsAcknowledgedNotification::new( + SubscriptionsAcknowledgedNotificationParams::new(accepted.as_ref().clone()), + ); + let mut meta = NotificationMetaObject::new(); + meta.set_subscription_id(request.id.clone()); + acknowledgment.extensions.insert(meta); + request + .peer + .send_notification(ServerNotification::SubscriptionsAcknowledgedNotification( + acknowledgment, + )) + .await + .map_err(|error| { + ErrorData::internal_error( + format!("failed to acknowledge subscription: {error}"), + None, + ) + })?; + Ok(Self { + request, + requested, + accepted, + sink, + _active_guard: active.drop_guard(), + }) + } + + /// Return the filter requested by the client. + pub fn requested(&self) -> &SubscriptionFilter { + &self.requested + } + + /// Return the subset accepted by the server. + pub fn accepted(&self) -> &SubscriptionFilter { + self.accepted.as_ref() + } + + /// Return a cloneable, filter-enforcing notification sink. + pub fn sink(&self) -> &SubscriptionSink { + &self.sink + } + + /// Wait until the subscription request is cancelled. + pub async fn cancelled(&self) { + self.request.ct.cancelled().await; + } + + /// Access the underlying request context. + pub fn request_context(&self) -> &RequestContext { + &self.request + } +} + impl> ServiceExt for S { fn serve_with_ct( self, @@ -161,20 +461,92 @@ where } } -/// Echoes the client-requested version if known; otherwise returns `server_fallback`. +/// Echoes the client-requested version if the server can serve it over the +/// `initialize` handshake; otherwise returns a legacy version the server does +/// support. +/// +/// `server_supported` comes from [`Service::supported_protocol_versions`], so a +/// server that narrows that list is never made to answer `initialize` with a +/// version it cannot serve. `2026-07-28` replaced the handshake with +/// per-request metadata, so a client naming that revision or later is answered +/// with the server's newest legacy version instead. +/// +/// # Errors +/// +/// Returns [`ErrorCode::UNSUPPORTED_PROTOCOL_VERSION`] when the server supports +/// no version that still has an `initialize` handshake. +/// +/// [`ErrorCode::UNSUPPORTED_PROTOCOL_VERSION`]: crate::model::ErrorCode::UNSUPPORTED_PROTOCOL_VERSION pub(crate) fn negotiate_protocol_version( client_requested: &ProtocolVersion, server_fallback: ProtocolVersion, -) -> ProtocolVersion { - if ProtocolVersion::KNOWN_VERSIONS.contains(client_requested) { - client_requested.clone() + server_supported: &[ProtocolVersion], +) -> Result { + if is_legacy_version(client_requested) && server_supported.contains(client_requested) { + return Ok(client_requested.clone()); + } + let legacy_fallback = if is_legacy_version(&server_fallback) { + Some(server_fallback) } else { + newest_legacy_version(server_supported) + }; + let Some(legacy_fallback) = legacy_fallback else { tracing::warn!( client_requested = %client_requested, - server_fallback = %server_fallback, - "client requested unsupported protocol version; falling back to server default" + "server supports no protocol version with an initialize handshake; rejecting" ); - server_fallback + return Err(ErrorData::unsupported_protocol_version( + client_requested.clone(), + server_supported, + )); + }; + // Falling back is the designed answer for a pinned client, and stateless + // HTTP re-runs it on every request, so this is not a warning. + tracing::debug!( + client_requested = %client_requested, + server_fallback = %legacy_fallback, + "client requested a protocol version unavailable over initialize; falling back to server default" + ); + Ok(legacy_fallback) +} + +/// The newest of `versions` that still has an `initialize` handshake. +fn newest_legacy_version(versions: &[ProtocolVersion]) -> Option { + versions + .iter() + .filter(|version| is_legacy_version(version)) + .max_by(|left, right| left.as_str().cmp(right.as_str())) + .cloned() +} + +fn missing_request_metadata_error(missing: &[&str]) -> ErrorData { + ErrorData::invalid_params( + format!( + "request _meta is missing or has malformed required fields: {}", + missing.join(", ") + ), + None, + ) +} + +/// Sends `error` as the response to the `initialize` request and reports it as +/// the reason the handshake failed. +async fn report_initialize_failure( + transport: &mut T, + error: ErrorData, + id: RequestId, +) -> ServerInitializeError +where + T: Transport + 'static, +{ + match transport + .send(ServerJsonRpcMessage::error(error.clone(), Some(id))) + .await + { + Ok(()) => ServerInitializeError::InitializeFailed(error), + Err(send_error) => { + ServerInitializeError::transport::(send_error, "sending error response") + } } } @@ -220,12 +592,48 @@ where } }; - let ClientRequest::InitializeRequest(peer_info) = &request else { - return Err(ServerInitializeError::ExpectedInitializeRequest(Some( - ClientJsonRpcMessage::request(request, id), - ))); + let initialize_request = match request { + ClientRequest::InitializeRequest(request) => request, + request => { + let missing_metadata = request + .get_meta() + .missing_required_keys(&ProtocolVersion::V_2026_07_28); + if !missing_metadata.is_empty() { + transport + .send(ServerJsonRpcMessage::error( + missing_request_metadata_error(&missing_metadata), + Some(id.clone()), + )) + .await + .map_err(|error| { + ServerInitializeError::transport::( + error, + "sending pre-init metadata error response", + ) + })?; + return Err(ServerInitializeError::ExpectedInitializeRequest(Some( + ClientJsonRpcMessage::request(request, id), + ))); + } + let (peer, peer_rx) = Peer::new(id_provider, None); + peer.require_request_metadata(); + // Dispatch the request from inside the service loop rather than + // inline: its handler may send notifications through `peer`, which + // only complete once the loop drains `peer_rx`. + return Ok(serve_inner( + service, + transport, + peer, + peer_rx, + VecDeque::from([ClientJsonRpcMessage::request(request, id)]), + ct, + )); + } }; - let (peer, peer_rx) = Peer::new(id_provider, Some(peer_info.params.clone())); + let requested_protocol_version = initialize_request.params.protocol_version.clone(); + let mut negotiated_peer_info = initialize_request.params.clone(); + let (peer, peer_rx) = Peer::new(id_provider, Some(negotiated_peer_info.clone())); + let request = ClientRequest::InitializeRequest(initialize_request); let context = RequestContext { ct: ct.child_token(), id: id.clone(), @@ -234,29 +642,23 @@ where peer: peer.clone(), }; // Send initialize response - let init_response = service.handle_request(request.clone(), context).await; - let mut init_response = match init_response { + let mut init_response = match service.handle_request(request, context).await { Ok(ServerResult::InitializeResult(init_response)) => init_response, Ok(result) => { return Err(ServerInitializeError::UnexpectedInitializeResponse(result)); } - Err(e) => { - transport - .send(ServerJsonRpcMessage::error(e.clone(), Some(id))) - .await - .map_err(|error| { - ServerInitializeError::transport::(error, "sending error response") - })?; - return Err(ServerInitializeError::InitializeFailed(e)); - } + Err(e) => return Err(report_initialize_failure(&mut transport, e, id).await), }; - init_response.protocol_version = negotiate_protocol_version( - &peer_info.params.protocol_version, + init_response.protocol_version = match negotiate_protocol_version( + &requested_protocol_version, init_response.protocol_version, - ); + &service.supported_protocol_versions(), + ) { + Ok(version) => version, + Err(e) => return Err(report_initialize_failure(&mut transport, e, id).await), + }; // Update peer_info so context.protocol_version() reflects the negotiated // version in all subsequent request handlers. - let mut negotiated_peer_info = peer_info.params.clone(); negotiated_peer_info.protocol_version = init_response.protocol_version.clone(); peer.set_peer_info(negotiated_peer_info); transport @@ -275,7 +677,14 @@ where // Streamable HTTP has no ordering guarantee between POSTs, and the MCP spec uses // SHOULD NOT (not MUST NOT) for pre-initialized messages, so any request arriving // before initialized is processed normally. - Ok(serve_inner(service, transport, peer, peer_rx, ct)) + Ok(serve_inner( + service, + transport, + peer, + peer_rx, + VecDeque::new(), + ct, + )) } macro_rules! method { @@ -430,6 +839,10 @@ impl Peer { } } + /// # SEP-2260: request association + /// + /// From protocol version `2026-07-28` this must be issued while handling a + /// client request; see [`OriginatingRequestId`]. #[deprecated( since = "1.8.0", note = "Sampling is deprecated by SEP-2577 and will be removed in a future release. See https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2577" @@ -464,6 +877,10 @@ impl Peer { } } method!( + /// # SEP-2260: request association + /// + /// From protocol version `2026-07-28` this must be issued while handling a + /// client request; see [`OriginatingRequestId`]. #[deprecated( since = "1.8.0", note = "Roots is deprecated by SEP-2577 and will be removed in a future release. See https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2577" @@ -471,11 +888,21 @@ impl Peer { peer_req list_roots ListRootsRequest() => ListRootsResult ); #[cfg(feature = "elicitation")] - method!(peer_req create_elicitation ElicitRequest(ElicitRequestParams) => ElicitResult); + method!( + /// # SEP-2260: request association + /// + /// From protocol version `2026-07-28` this must be issued while handling a + /// client request; see [`OriginatingRequestId`]. + peer_req create_elicitation ElicitRequest(ElicitRequestParams) => ElicitResult + ); #[cfg(feature = "elicitation")] - method!(peer_req_with_timeout create_elicitation_with_timeout ElicitRequest(ElicitRequestParams) => ElicitResult); - #[cfg(feature = "elicitation")] - method!(peer_not notify_url_elicitation_completed ElicitationCompleteNotification(ElicitationResponseNotificationParam)); + method!( + /// # SEP-2260: request association + /// + /// From protocol version `2026-07-28` this must be issued while handling a + /// client request; see [`OriginatingRequestId`]. + peer_req_with_timeout create_elicitation_with_timeout ElicitRequest(ElicitRequestParams) => ElicitResult + ); method!(peer_not notify_cancelled CancelledNotification(CancelledNotificationParam)); method!(peer_not notify_progress ProgressNotification(ProgressNotificationParam)); @@ -702,6 +1129,11 @@ impl Peer { /// # Ok(()) /// # } /// ``` + /// + /// # SEP-2260: request association + /// + /// From protocol version `2026-07-28` this must be issued while handling a + /// client request; see [`OriginatingRequestId`]. #[cfg(all(feature = "schemars", feature = "elicitation"))] pub async fn elicit(&self, message: impl Into) -> Result, ElicitationError> where @@ -763,6 +1195,11 @@ impl Peer { /// # Ok(()) /// # } /// ``` + /// + /// # SEP-2260: request association + /// + /// From protocol version `2026-07-28` this must be issued while handling a + /// client request; see [`OriginatingRequestId`]. #[cfg(all(feature = "schemars", feature = "elicitation"))] pub async fn elicit_with_timeout( &self, @@ -858,6 +1295,11 @@ impl Peer { /// Ok(()) /// } /// ``` + /// + /// # SEP-2260: request association + /// + /// From protocol version `2026-07-28` this must be issued while handling a + /// client request; see [`OriginatingRequestId`]. #[cfg(feature = "elicitation")] pub async fn elicit_url( &self, @@ -909,6 +1351,11 @@ impl Peer { /// Ok(()) /// } /// ``` + /// + /// # SEP-2260: request association + /// + /// From protocol version `2026-07-28` this must be issued while handling a + /// client request; see [`OriginatingRequestId`]. #[cfg(feature = "elicitation")] pub async fn elicit_url_with_timeout( &self, diff --git a/crates/rmcp/src/task_manager.rs b/crates/rmcp/src/task_manager.rs index 32bcf8f..0e8339c 100644 --- a/crates/rmcp/src/task_manager.rs +++ b/crates/rmcp/src/task_manager.rs @@ -1,302 +1,998 @@ -use std::{any::Any, collections::HashMap, pin::Pin}; +//! Server-side runtime for the MCP Tasks extension (SEP-2663, +//! `io.modelcontextprotocol/tasks`). +//! +//! [`TaskManager`] owns the durable state for tasks a server has materialized +//! in response to task-eligible requests (currently `tools/call`). It: +//! +//! - spawns the underlying operation and tracks its lifecycle as a +//! [`DetailedTask`] (`working` → terminal, optionally via `input_required`), +//! - answers `tasks/get` with the current state (including in-flight +//! `inputRequests` and terminal `result`/`error` payloads), +//! - accepts `tasks/update` `inputResponses` and routes them to the running +//! operation (ignoring unknown or already-answered keys per spec), +//! - handles cooperative `tasks/cancel`, +//! - enforces TTL-based expiry (`ttl_ms`), marking overdue tasks `failed`. +//! +//! Tasks are only durably observable once [`TaskManager::spawn`] returns, +//! satisfying the spec requirement that a server not return `CreateTaskResult` +//! before `tasks/get` for that id would resolve. + +use std::{ + collections::HashMap, + pin::Pin, + sync::{Arc, Mutex}, + time::Instant, +}; use futures::Future; -use tokio::{ - sync::mpsc, - time::{Duration, timeout}, -}; +use tokio::sync::oneshot; use crate::{ - RoleServer, - error::{ErrorData as McpError, RmcpError as Error}, - model::{CallToolResult, ClientRequest}, - service::RequestContext, + error::ErrorData as McpError, + model::{ + CallToolResult, DetailedTask, InputRequest, InputRequests, JsonObject, Task, TaskPayload, + TaskStatus, + }, }; -/// Boxed future that represents an asynchronous operation managed by the processor. -pub type OperationFuture = - Pin, Error>> + Send>>; +/// Default TTL (5 minutes, in milliseconds) applied when none is specified. +pub const DEFAULT_TASK_TTL_MS: u64 = 300_000; -/// Describes metadata associated with an enqueued task. -#[derive(Debug, Clone)] -#[non_exhaustive] -pub struct OperationDescriptor { - pub operation_id: String, - pub name: String, - pub client_request: Option, - pub context: Option>, - pub ttl: Option, -} - -impl OperationDescriptor { - pub fn new(operation_id: impl Into, name: impl Into) -> Self { - Self { - operation_id: operation_id.into(), - name: name.into(), - client_request: None, - context: None, - ttl: None, - } - } - - pub fn with_client_request(mut self, request: ClientRequest) -> Self { - self.client_request = Some(request); - self - } - - pub fn with_context(mut self, context: RequestContext) -> Self { - self.context = Some(context); - self - } - - pub fn with_ttl(mut self, ttl: u64) -> Self { - self.ttl = Some(ttl); - self - } -} - -/// Operation message describing a unit of asynchronous work. -#[non_exhaustive] -pub struct OperationMessage { - pub descriptor: OperationDescriptor, - pub future: OperationFuture, -} - -impl OperationMessage { - pub fn new(descriptor: OperationDescriptor, future: OperationFuture) -> Self { - Self { descriptor, future } - } -} - -/// Trait for operation result transport -pub trait OperationResultTransport: Send + Sync + 'static { - fn operation_id(&self) -> &String; - fn as_any(&self) -> &dyn std::any::Any; -} - -// ===== Operation Processor ===== -pub const DEFAULT_TASK_TIMEOUT_SECS: u64 = 300; // 5 minutes -/// Operation processor that coordinates extractors and handlers -pub struct OperationProcessor { - /// Currently running tasks keyed by id - running_tasks: HashMap, - /// Completed results waiting to be collected - completed_results: Vec, - task_result_receiver: mpsc::UnboundedReceiver, - task_result_sender: mpsc::UnboundedSender, -} - -struct RunningTask { - task_handle: tokio::task::JoinHandle<()>, - started_at: std::time::Instant, - timeout: Option, - descriptor: OperationDescriptor, -} - -#[non_exhaustive] -pub struct TaskResult { - pub descriptor: OperationDescriptor, - pub result: Result, Error>, -} +/// Default suggested polling interval, in milliseconds. +pub const DEFAULT_POLL_INTERVAL_MS: u64 = 1_000; /// Helper to generate an ISO 8601 timestamp for task metadata. pub fn current_timestamp() -> String { chrono::Utc::now().to_rfc3339() } -/// Result transport for tool calls executed as tasks. -pub struct ToolCallTaskResult { - id: String, - pub result: Result, +/// Handle passed to a running task operation, allowing it to surface +/// server-to-client requests (elicitation, sampling, roots) mid-task and +/// await the client's `tasks/update` response. +#[derive(Clone)] +pub struct TaskContext { + task_id: String, + inner: Arc>, } -impl ToolCallTaskResult { - pub fn new(id: impl Into, result: Result) -> Self { - Self { - id: id.into(), - result, +impl TaskContext { + /// The id of the task this context belongs to. + pub fn task_id(&self) -> &str { + &self.task_id + } + + /// Surface a server-to-client request under `key` and wait for the + /// client's response delivered via `tasks/update`. + /// + /// While at least one request is outstanding the task reports + /// `input_required` from `tasks/get`, with all outstanding requests in + /// `inputRequests`. Keys must be unique over the lifetime of the task; + /// reusing a key returns an error. + pub async fn request_input( + &self, + key: impl Into, + request: InputRequest, + ) -> Result { + let key = key.into(); + let (tx, rx) = oneshot::channel(); + { + let mut inner = self.inner.lock().expect("task manager lock poisoned"); + let entry = inner.tasks.get_mut(&self.task_id).ok_or_else(|| { + TaskExit::Error(McpError::internal_error( + "task no longer exists".to_string(), + None, + )) + })?; + if !entry.used_input_keys.insert(key.clone()) { + return Err(TaskExit::Error(McpError::internal_error( + format!("inputRequests key {key:?} was already used for this task"), + None, + ))); + } + entry.pending_inputs.insert(key.clone(), (request, tx)); + entry.touch(); + } + // The sender is dropped when `tasks/cancel` clears pending inputs. + rx.await.map_err(|_| TaskExit::Cancelled) + } + + /// Update the task's human-readable status message. + pub fn set_status_message(&self, message: impl Into) { + let mut inner = self.inner.lock().expect("task manager lock poisoned"); + if let Some(entry) = inner.tasks.get_mut(&self.task_id) { + entry.task.status_message = Some(message.into()); + entry.touch(); + } + } + + /// Returns `true` if `tasks/cancel` has been received for this task. + /// Cooperative: operations should check this and stop when set. + pub fn is_cancel_requested(&self) -> bool { + let inner = self.inner.lock().expect("task manager lock poisoned"); + inner + .tasks + .get(&self.task_id) + .is_some_and(|e| e.cancel_requested) + } + + /// Resolves once `tasks/cancel` has been received for this task (or + /// immediately, if it already has). Cooperative: pair with + /// `tokio::select!` around long-running work to implement a cancellation + /// exit path. + /// + /// An operation that stops in response should return + /// [`TaskExit::Cancelled`] so the task settles as `cancelled`. Returning + /// [`TaskExit::Error`] settles as `failed`, and finishing the work + /// anyway settles as `completed` — per SEP-2663 cancellation is + /// cooperative and a task may reach a non-`cancelled` terminal status. + pub async fn cancelled(&self) { + let mut rx = { + let inner = self.inner.lock().expect("task manager lock poisoned"); + let Some(entry) = inner.tasks.get(&self.task_id) else { + return; + }; + if entry.cancel_requested { + return; + } + entry.cancel_signal.subscribe() + }; + // Wait until the watch flips to true; a closed channel means the + // manager dropped the entry, which also unblocks the operation. + while !*rx.borrow_and_update() { + if rx.changed().await.is_err() { + return; + } } } } -impl OperationResultTransport for ToolCallTaskResult { - fn operation_id(&self) -> &String { - &self.id +/// How a task operation finished without producing a result. +#[expect( + clippy::exhaustive_enums, + reason = "error variant for task exit may only be due to error or cancellation" +)] +#[derive(Debug)] +pub enum TaskExit { + /// The operation is exiting in response to a cancellation request; + /// the task settles as terminal `cancelled`. + Cancelled, + /// A real failure; the task settles as terminal `failed` with the + /// error inlined, even after `tasks/cancel` was received. + Error(McpError), +} + +impl From for TaskExit { + fn from(error: McpError) -> Self { + TaskExit::Error(error) + } +} + +/// Boxed future representing the async operation backing a task. +pub type TaskFuture = Pin> + Send>>; + +struct TaskEntry { + task: Task, + /// Terminal payload, if the task has finished. + terminal: Option, + /// When the task reached its terminal state; drives retention eviction. + terminal_at: Option, + /// Outstanding input requests keyed by their unique identifier. + pending_inputs: HashMap)>, + /// Every key ever used, to enforce uniqueness across the task lifetime. + used_input_keys: std::collections::HashSet, + cancel_requested: bool, + /// Signals the running operation that cancellation was requested + /// (`true` once `tasks/cancel` arrives). Cooperative: the operation + /// decides whether and how to stop. + cancel_signal: tokio::sync::watch::Sender, + created: Instant, + join_handle: Option>, +} + +impl TaskEntry { + fn touch(&mut self) { + self.task.last_updated_at = current_timestamp(); } - fn as_any(&self) -> &dyn Any { + fn current_status(&self) -> TaskStatus { + match &self.terminal { + Some(payload) => payload.status(), + None if !self.pending_inputs.is_empty() => TaskStatus::InputRequired, + None => TaskStatus::Working, + } + } + + fn detailed(&self) -> DetailedTask { + let payload = match &self.terminal { + Some(p) => p.clone(), + None if !self.pending_inputs.is_empty() => TaskPayload::InputRequired { + input_requests: self + .pending_inputs + .iter() + .map(|(k, (req, _))| (k.clone(), req.clone())) + .collect::(), + }, + None => TaskPayload::Working, + }; + DetailedTask::new(self.task.clone(), payload) + } +} + +#[derive(Default)] +struct TaskManagerInner { + tasks: HashMap, +} + +/// Options controlling a spawned task. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct TaskOptions { + /// TTL in milliseconds; `None` means unlimited retention. + pub ttl_ms: Option, + /// Suggested polling interval in milliseconds. + pub poll_interval_ms: Option, + /// Initial status message. + pub status_message: Option, +} + +impl Default for TaskOptions { + fn default() -> Self { + Self { + ttl_ms: Some(DEFAULT_TASK_TTL_MS), + poll_interval_ms: Some(DEFAULT_POLL_INTERVAL_MS), + status_message: None, + } + } +} + +impl TaskOptions { + pub fn new() -> Self { + Self::default() + } + + /// Set the TTL in milliseconds. `None` means unlimited retention. + pub fn with_ttl_ms(mut self, ttl_ms: impl Into>) -> Self { + self.ttl_ms = ttl_ms.into(); + self + } + + /// Set the suggested polling interval in milliseconds. + pub fn with_poll_interval_ms(mut self, poll_interval_ms: u64) -> Self { + self.poll_interval_ms = Some(poll_interval_ms); + self + } + + /// Set the initial status message. + pub fn with_status_message(mut self, message: impl Into) -> Self { + self.status_message = Some(message.into()); self } } -impl Default for OperationProcessor { - fn default() -> Self { - Self::new() - } +/// Server-side task store and executor for the SEP-2663 Tasks extension. +/// +/// Cheaply cloneable; all clones share the same state. +/// +/// # Retention +/// +/// Entries are swept opportunistically on every `spawn` / `get_task` / +/// `update_task` / `cancel_task` / `running_task_count` call: non-terminal +/// tasks whose `ttl_ms` has elapsed are marked `failed` (their operation is +/// aborted), and terminal tasks are evicted after being retained for one +/// further `ttl_ms` window past their terminal transition so pollers can +/// observe the final state. +/// +/// Note that the retention window intentionally extends past the +/// creation-based lifetime that `ttl_ms` advertises on the wire: a task that +/// runs to its TTL deadline is marked `failed` around `created + ttl_ms` and +/// stays observable until roughly `created + 2 × ttl_ms`. This is compliant — +/// SEP-2663 lets servers delete expired tasks *at any time* after the TTL, +/// so retaining them longer as an observation grace period is a server-side +/// policy choice, not a wire-contract change. Clients may treat the task as +/// unusable after `createdAt + ttlMs` regardless. +/// +/// Tasks with `ttl_ms: None` are retained for the lifetime of the manager +/// (spec: unlimited retention) — bound task creation or call +/// [`Self::shutdown`] yourself if you spawn such tasks in a long-lived +/// server. There is no background sweeper; an idle manager holds its entries +/// until the next call. +#[derive(Clone, Default)] +pub struct TaskManager { + inner: Arc>, } -impl OperationProcessor { +impl TaskManager { pub fn new() -> Self { - let (task_result_sender, task_result_receiver) = mpsc::unbounded_channel(); - Self { - running_tasks: HashMap::new(), - completed_results: Vec::new(), - task_result_receiver, - task_result_sender, - } + Self::default() } - /// Submit an operation for asynchronous execution. - #[allow(clippy::result_large_err)] - pub fn submit_operation(&mut self, message: OperationMessage) -> Result<(), Error> { - if self - .running_tasks - .contains_key(&message.descriptor.operation_id) + /// Spawn an operation as a task and return its seed [`Task`] state for a + /// `CreateTaskResult`. The task is durably observable via + /// [`Self::get_task`] before this method returns. + /// + /// `make_future` receives a [`TaskContext`] for mid-task input requests, + /// status messages, and cooperative cancellation checks. + pub fn spawn(&self, options: TaskOptions, make_future: F) -> Task + where + F: FnOnce(TaskContext) -> TaskFuture, + { + let task_id = uuid::Uuid::new_v4().to_string(); + let now = current_timestamp(); + let mut task = Task::new(task_id.clone(), TaskStatus::Working, now.clone(), now); + task.ttl_ms = options.ttl_ms; + task.poll_interval_ms = options.poll_interval_ms; + task.status_message = options.status_message; + + let entry = TaskEntry { + task: task.clone(), + terminal: None, + terminal_at: None, + pending_inputs: HashMap::new(), + used_input_keys: std::collections::HashSet::new(), + cancel_requested: false, + cancel_signal: tokio::sync::watch::channel(false).0, + created: Instant::now(), + join_handle: None, + }; { - return Err(Error::TaskError(format!( - "Operation with id {} is already running", - message.descriptor.operation_id - ))); + let mut inner = self.inner.lock().expect("task manager lock poisoned"); + // Opportunistic TTL sweep on every task creation, so terminal + // entries are evicted even if clients never poll again. + Self::sweep_expired(&mut inner); + inner.tasks.insert(task_id.clone(), entry); } - self.spawn_async_task(message); + + let context = TaskContext { + task_id: task_id.clone(), + inner: self.inner.clone(), + }; + let future = make_future(context); + let inner = self.inner.clone(); + let id_for_task = task_id.clone(); + let originating_request = crate::service::ORIGINATING_REQUEST + .try_with(|id| id.clone()) + .ok(); + let handle = tokio::spawn(async move { + let result = run_task_operation(originating_request, future).await; + let mut inner = inner.lock().expect("task manager lock poisoned"); + if let Some(entry) = inner.tasks.get_mut(&id_for_task) { + if entry.terminal.is_none() { + entry.terminal = Some(match result { + Ok(result) => TaskPayload::Completed { + result: result_to_object(&result), + }, + Err(TaskExit::Cancelled) => TaskPayload::Cancelled, + Err(TaskExit::Error(error)) => TaskPayload::Failed { + error: error_to_object(&error), + }, + }); + entry.terminal_at = Some(Instant::now()); + entry.pending_inputs.clear(); + entry.touch(); + entry.task.status = entry.current_status(); + } + // The operation has finished; drop the JoinHandle so it is + // not retained for the rest of the retention window. + entry.join_handle = None; + } + }); + match self + .inner + .lock() + .expect("task manager lock poisoned") + .tasks + .get_mut(&task_id) + { + // Only store the handle while the operation is still running: if + // it already settled, the completion path above ran first and a + // stored handle would never be cleared. + Some(entry) => { + if entry.terminal.is_none() { + entry.join_handle = Some(handle); + } + } + // The entry is gone: shutdown() drained the map between the + // insert and here. Abort rather than leak a detached operation. + None => handle.abort(), + } + task + } + + /// Handle `tasks/get`: return the current [`DetailedTask`] state. + pub fn get_task(&self, task_id: &str) -> Result { + let mut inner = self.inner.lock().expect("task manager lock poisoned"); + Self::sweep_expired(&mut inner); + let entry = inner + .tasks + .get_mut(task_id) + .ok_or_else(|| unknown_task(task_id))?; + entry.task.status = entry.current_status(); + Ok(entry.detailed()) + } + + /// Handle `tasks/update`: deliver `inputResponses` to the running + /// operation. Unknown, already-answered, or superseded keys are ignored + /// per spec; a partial set of responses is accepted. + pub fn update_task( + &self, + task_id: &str, + input_responses: impl IntoIterator, + ) -> Result<(), McpError> { + let mut inner = self.inner.lock().expect("task manager lock poisoned"); + Self::sweep_expired(&mut inner); + let entry = inner + .tasks + .get_mut(task_id) + .ok_or_else(|| unknown_task(task_id))?; + for (key, value) in input_responses { + if let Some((_, tx)) = entry.pending_inputs.remove(&key) { + // Receiver dropped means the operation moved on; ignore. + let _ = tx.send(value); + } + } + entry.touch(); + entry.task.status = entry.current_status(); Ok(()) } - fn spawn_async_task(&mut self, message: OperationMessage) { - let OperationMessage { descriptor, future } = message; - let task_id = descriptor.operation_id.clone(); - let timeout_secs = descriptor.ttl.or(Some(DEFAULT_TASK_TIMEOUT_SECS)); - let sender = self.task_result_sender.clone(); - let descriptor_for_result = descriptor.clone(); + /// Handle `tasks/cancel`: cooperative cancellation (SEP-2663). + /// + /// Records the cancellation *intent* and acknowledges immediately, but + /// does **not** abort the underlying future or force a terminal state. + /// The operation observes cancellation via + /// [`TaskContext::is_cancel_requested`] / [`TaskContext::cancelled`], or + /// via the error returned from a pending [`TaskContext::request_input`] + /// call (whose response channel is dropped here), and decides its own + /// terminal status: + /// + /// - stops with [`TaskExit::Cancelled`] → recorded as `cancelled`, + /// - stops with [`TaskExit::Error`] → recorded as `failed` with the + /// error inlined (a real failure after a cancel request is not masked), + /// - finishes its work anyway → recorded as `completed` — per the spec, + /// "the task may still reach a non-`cancelled` terminal status". + pub fn cancel_task(&self, task_id: &str) -> Result<(), McpError> { + let mut inner = self.inner.lock().expect("task manager lock poisoned"); + Self::sweep_expired(&mut inner); + let entry = inner + .tasks + .get_mut(task_id) + .ok_or_else(|| unknown_task(task_id))?; + entry.cancel_requested = true; + let _ = entry.cancel_signal.send(true); + if entry.terminal.is_none() { + // Wake any operation parked on `request_input`: dropping the + // response senders resolves those awaits with an error, giving + // parked operations a cooperative exit path. The task leaves + // `input_required` and reports `working` until it settles. + entry.pending_inputs.clear(); + entry.touch(); + entry.task.status = entry.current_status(); + } + Ok(()) + } - let timed_future = async move { - if let Some(secs) = timeout_secs { - match timeout(Duration::from_secs(secs), future).await { - Ok(result) => result, - Err(_) => Err(Error::TaskError("Operation timed out".to_string())), - } - } else { - future.await + /// Number of tasks currently in a non-terminal state. + pub fn running_task_count(&self) -> usize { + let mut inner = self.inner.lock().expect("task manager lock poisoned"); + Self::sweep_expired(&mut inner); + inner + .tasks + .values() + .filter(|e| e.terminal.is_none()) + .count() + } + + /// Abort all running tasks and clear all task state. + pub fn shutdown(&self) { + let mut inner = self.inner.lock().expect("task manager lock poisoned"); + for (_, mut entry) in inner.tasks.drain() { + if let Some(handle) = entry.join_handle.take() { + handle.abort(); } - }; + } + } - let handle = tokio::spawn(async move { - let result = timed_future.await; - let task_result = TaskResult { - descriptor: descriptor_for_result, - result, + /// TTL sweep, run from every `TaskManager` entry point (SEP-2663: servers + /// MAY mark a task `failed` any time after its TTL elapses, and + /// subsequently delete it at any time; `ttl_ms: None` means unlimited + /// retention). + /// + /// Two phases per entry: + /// 1. A non-terminal task whose TTL has elapsed is marked `failed` (its + /// operation is aborted — the TTL is the SDK's hard-stop safety valve, + /// unlike cooperative `tasks/cancel`). + /// 2. A *terminal* task is evicted once it has been retained for a full + /// TTL window after reaching its terminal state, so well-behaved + /// pollers get a chance to observe the final status before late + /// `tasks/get` calls start returning `-32602`. + fn sweep_expired(inner: &mut TaskManagerInner) { + // Phase 1: fail overdue non-terminal tasks. + for entry in inner.tasks.values_mut() { + if entry.terminal.is_none() + && let Some(ttl_ms) = entry.task.ttl_ms + && entry.created.elapsed().as_millis() >= u128::from(ttl_ms) + { + if let Some(handle) = entry.join_handle.take() { + handle.abort(); + } + entry.terminal = Some(TaskPayload::Failed { + error: error_to_object(&McpError::internal_error( + "task expired: TTL elapsed before completion".to_string(), + None, + )), + }); + entry.terminal_at = Some(Instant::now()); + entry.pending_inputs.clear(); + entry.touch(); + entry.task.status = TaskStatus::Failed; + } + } + // Phase 2: evict terminal tasks whose retention window has passed. + inner.tasks.retain(|_, entry| { + let (Some(ttl_ms), Some(terminal_at)) = (entry.task.ttl_ms, entry.terminal_at) else { + return true; }; - let _ = sender.send(task_result); + terminal_at.elapsed().as_millis() < u128::from(ttl_ms) }); - let running_task = RunningTask { - task_handle: handle, - started_at: std::time::Instant::now(), - timeout: timeout_secs, - descriptor, - }; - self.running_tasks.insert(task_id, running_task); + } +} + +fn unknown_task(task_id: &str) -> McpError { + McpError::invalid_params(format!("unknown task: {task_id}"), None) +} + +async fn run_task_operation( + originating_request: Option, + future: TaskFuture, +) -> Result { + match originating_request { + Some(id) => crate::service::ORIGINATING_REQUEST.scope(id, future).await, + None => future.await, + } +} + +fn result_to_object(result: &CallToolResult) -> JsonObject { + match serde_json::to_value(result) { + Ok(serde_json::Value::Object(map)) => map, + _ => JsonObject::new(), + } +} + +fn error_to_object(error: &McpError) -> JsonObject { + match serde_json::to_value(error) { + Ok(serde_json::Value::Object(map)) => map, + _ => JsonObject::new(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::model::ContentBlock; + + fn ok_result(text: &str) -> CallToolResult { + CallToolResult::success(vec![ContentBlock::text(text.to_string())]) + } + + #[tokio::test] + async fn task_completes_and_result_is_inlined() { + let manager = TaskManager::new(); + let task = manager.spawn(TaskOptions::default(), |_ctx| { + Box::pin(async { Ok(ok_result("42")) }) + }); + assert_eq!(task.status, TaskStatus::Working); + + // Durable immediately. + manager.get_task(&task.task_id).unwrap(); + + // Wait for completion. + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + let detailed = manager.get_task(&task.task_id).unwrap(); + if detailed.status() == TaskStatus::Completed { + match detailed.payload { + TaskPayload::Completed { result } => { + assert!(result.contains_key("content")); + return; + } + other => panic!("unexpected payload: {other:?}"), + } + } + } + panic!("task did not complete"); } + + #[tokio::test] + async fn cancel_settles_to_cancelled_when_operation_honors_it() { + let manager = TaskManager::new(); + let task = manager.spawn(TaskOptions::default(), |ctx| { + Box::pin(async move { + tokio::select! { + _ = ctx.cancelled() => Err(TaskExit::Cancelled), + _ = tokio::time::sleep(std::time::Duration::from_secs(60)) => { + Ok(ok_result("never")) + } + } + }) + }); + manager.cancel_task(&task.task_id).unwrap(); - /// Collect completed results from running tasks and remove them from the running tasks map. - fn collect_completed_results(&mut self) { - while let Ok(result) = self.task_result_receiver.try_recv() { - self.running_tasks.remove(&result.descriptor.operation_id); - self.completed_results.push(result); + // The ack is immediate but the terminal state is set by the + // operation; poll until it settles. + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + let detailed = manager.get_task(&task.task_id).unwrap(); + if detailed.status().is_terminal() { + assert_eq!(detailed.status(), TaskStatus::Cancelled); + return; + } } + panic!("task did not settle after cancel"); } - /// Check for tasks that have exceeded their timeout and handle them appropriately. - pub fn check_timeouts(&mut self) { - self.collect_completed_results(); - let now = std::time::Instant::now(); - let mut timed_out_tasks = Vec::new(); + #[tokio::test] + async fn post_cancel_unrelated_error_settles_as_failed() { + let manager = TaskManager::new(); + let task = manager.spawn(TaskOptions::default(), |ctx| { + Box::pin(async move { + // Fail for an unrelated reason after observing the cancel: + // must be recorded as `failed`, not masked as `cancelled`. + ctx.cancelled().await; + Err(TaskExit::Error(McpError::internal_error( + "database write failed", + None, + ))) + }) + }); + manager.cancel_task(&task.task_id).unwrap(); - for (task_id, task) in &self.running_tasks { - if let Some(timeout_duration) = task.timeout { - if now.duration_since(task.started_at).as_secs() > timeout_duration { - task.task_handle.abort(); - timed_out_tasks.push(task_id.clone()); + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + let detailed = manager.get_task(&task.task_id).unwrap(); + if detailed.status().is_terminal() { + assert_eq!(detailed.status(), TaskStatus::Failed); + match detailed.payload { + TaskPayload::Failed { error } => { + assert!( + error.get("message").is_some_and(|m| m + .as_str() + .is_some_and(|s| s.contains("database write failed"))), + "error payload should be preserved: {error:?}" + ); + } + other => panic!("unexpected payload: {other:?}"), } + return; } } + panic!("task did not settle after cancel"); + } + + #[tokio::test] + async fn cancel_is_cooperative_and_lets_the_operation_clean_up() { + let manager = TaskManager::new(); + let (cleanup_tx, cleanup_rx) = oneshot::channel::<&'static str>(); + let task = manager.spawn(TaskOptions::default(), |ctx| { + Box::pin(async move { + // Wait for cancellation, then run cleanup and finish the + // work anyway (spec: a task may still reach a + // non-`cancelled` terminal status). + ctx.cancelled().await; + let _ = cleanup_tx.send("cleaned up"); + Ok(ok_result("finished despite cancel")) + }) + }); + + manager.cancel_task(&task.task_id).unwrap(); - for task_id in timed_out_tasks { - if let Some(task) = self.running_tasks.remove(&task_id) { - let timeout_result = TaskResult { - descriptor: task.descriptor, - result: Err(Error::TaskError("Operation timed out".to_string())), - }; - self.completed_results.push(timeout_result); + // The ack is immediate and does not force a terminal state. + let detailed = manager.get_task(&task.task_id).unwrap(); + assert!( + !detailed.status().is_terminal(), + "cancel must not force terminal state" + ); + + // The operation observes the cancel and performs cleanup. + let cleanup = tokio::time::timeout(std::time::Duration::from_secs(5), cleanup_rx) + .await + .expect("cleanup should not time out") + .expect("cleanup channel should not be dropped"); + assert_eq!(cleanup, "cleaned up"); + + // The operation chose to complete: the task settles as `completed`. + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + let detailed = manager.get_task(&task.task_id).unwrap(); + if detailed.status().is_terminal() { + assert_eq!(detailed.status(), TaskStatus::Completed); + return; } } + panic!("task did not settle after cancel"); } - /// Get the number of running tasks. - pub fn running_task_count(&mut self) -> usize { - self.collect_completed_results(); - self.running_tasks.len() + #[tokio::test] + async fn cancel_wakes_parked_input_requests() { + let manager = TaskManager::new(); + let (exit_tx, exit_rx) = oneshot::channel::<&'static str>(); + let task = manager.spawn(TaskOptions::default(), |ctx| { + Box::pin(async move { + let request: InputRequest = serde_json::from_value(serde_json::json!({ + "method": "elicitation/create", + "params": { + "message": "Waiting forever", + "requestedSchema": {"type": "object", "properties": {}} + } + })) + .map_err(|e| McpError::internal_error(e.to_string(), None))?; + // Parked on input; cancel must wake this await with an error. + let err = ctx.request_input("k1", request).await.unwrap_err(); + let _ = exit_tx.send("woken"); + Err(err) + }) + }); + + // Wait until the task is parked on the input request. + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + if manager.get_task(&task.task_id).unwrap().status() == TaskStatus::InputRequired { + break; + } + } + + manager.cancel_task(&task.task_id).unwrap(); + let woken = tokio::time::timeout(std::time::Duration::from_secs(5), exit_rx) + .await + .expect("parked operation should be woken by cancel") + .expect("exit channel should not be dropped"); + assert_eq!(woken, "woken"); + assert_eq!( + manager.get_task(&task.task_id).unwrap().status(), + TaskStatus::Cancelled + ); } - /// Cancel all running tasks. - pub fn cancel_all_tasks(&mut self) { - for (_, task) in self.running_tasks.drain() { - task.task_handle.abort(); + #[tokio::test] + async fn unknown_task_is_invalid_params() { + // SEP-2663 §Protocol Errors: invalid or nonexistent taskId is -32602 + // (Invalid params) — MUST for tasks/get, SHOULD for update/cancel. + let manager = TaskManager::new(); + for err in [ + manager.get_task("nope").unwrap_err(), + manager.cancel_task("nope").unwrap_err(), + manager.update_task("nope", []).unwrap_err(), + ] { + assert_eq!(err.code, crate::model::ErrorCode::INVALID_PARAMS); } - while self.task_result_receiver.try_recv().is_ok() {} - self.completed_results.clear(); } + + #[tokio::test] + async fn terminal_tasks_are_evicted_after_retention_window() { + let manager = TaskManager::new(); + let task = manager.spawn(TaskOptions::new().with_ttl_ms(50), |_ctx| { + Box::pin(async { Ok(ok_result("fast")) }) + }); - /// List running task ids. - pub fn list_running(&mut self) -> Vec { - self.collect_completed_results(); - self.running_tasks.keys().cloned().collect() + // Wait for completion; the terminal state stays observable during + // the retention window. + let mut completed = false; + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + if manager.get_task(&task.task_id).unwrap().status() == TaskStatus::Completed { + completed = true; + break; + } + } + assert!(completed, "task should have completed"); + + // After a full TTL window past terminal, the entry is evicted and + // late polls get -32602. + tokio::time::sleep(std::time::Duration::from_millis(120)).await; + let err = manager.get_task(&task.task_id).unwrap_err(); + assert_eq!(err.code, crate::model::ErrorCode::INVALID_PARAMS); + assert_eq!(manager.running_task_count(), 0); + } + + #[tokio::test] + async fn abandoned_tasks_are_swept_by_other_entry_points() { + // A task nobody ever polls again must still be failed + evicted; the + // sweep runs from spawn() too, so activity on *other* tasks is enough. + let manager = TaskManager::new(); + let abandoned = manager.spawn(TaskOptions::new().with_ttl_ms(10), |_ctx| { + Box::pin(async { + tokio::time::sleep(std::time::Duration::from_secs(60)).await; + Ok(ok_result("never")) + }) + }); + + // Let TTL elapse (fails the task), then a second full window + // (evicts it), without ever calling get_task on the abandoned id. + tokio::time::sleep(std::time::Duration::from_millis(40)).await; + let _ = manager.spawn(TaskOptions::default(), |_ctx| { + Box::pin(async { Ok(ok_result("other")) }) + }); + tokio::time::sleep(std::time::Duration::from_millis(40)).await; + let _ = manager.spawn(TaskOptions::default(), |_ctx| { + Box::pin(async { Ok(ok_result("other2")) }) + }); + + let err = manager.get_task(&abandoned.task_id).unwrap_err(); + assert_eq!(err.code, crate::model::ErrorCode::INVALID_PARAMS); } - /// Returns a snapshot of completed task results. - pub fn peek_completed(&mut self) -> &[TaskResult] { - self.collect_completed_results(); - &self.completed_results + #[tokio::test] + async fn running_task_count_sweeps_expired_tasks() { + let manager = TaskManager::new(); + let _task = manager.spawn(TaskOptions::new().with_ttl_ms(10), |_ctx| { + Box::pin(async { + tokio::time::sleep(std::time::Duration::from_secs(60)).await; + Ok(ok_result("never")) + }) + }); + assert_eq!(manager.running_task_count(), 1); + tokio::time::sleep(std::time::Duration::from_millis(30)).await; + // The count itself must sweep: the overdue task is failed, not + // reported as running. + assert_eq!(manager.running_task_count(), 0); } + + #[tokio::test] + async fn unlimited_ttl_tasks_are_retained() { + let manager = TaskManager::new(); + let task = manager.spawn(TaskOptions::new().with_ttl_ms(None), |_ctx| { + Box::pin(async { Ok(ok_result("kept")) }) + }); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + // Sweeps triggered by other entry points must not evict it. + let _ = manager.spawn(TaskOptions::default(), |_ctx| { + Box::pin(async { Ok(ok_result("other")) }) + }); + assert_eq!( + manager.get_task(&task.task_id).unwrap().status(), + TaskStatus::Completed + ); + } + + #[tokio::test] + async fn ttl_expiry_fails_task() { + let manager = TaskManager::new(); + let task = manager.spawn( + TaskOptions { + ttl_ms: Some(10), + ..Default::default() + }, + |_ctx| { + Box::pin(async { + tokio::time::sleep(std::time::Duration::from_secs(60)).await; + Ok(ok_result("never")) + }) + }, + ); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + let detailed = manager.get_task(&task.task_id).unwrap(); + assert_eq!(detailed.status(), TaskStatus::Failed); + } + + #[tokio::test] + async fn input_required_roundtrip() { + let manager = TaskManager::new(); + let task = manager.spawn(TaskOptions::default(), |ctx| { + Box::pin(async move { + let request: InputRequest = serde_json::from_value(serde_json::json!({ + "method": "elicitation/create", + "params": { + "message": "What is your name?", + "requestedSchema": {"type": "object", "properties": {}} + } + })) + .map_err(|e| McpError::internal_error(e.to_string(), None))?; + let response = ctx.request_input("name-1", request).await?; + let name = response + .get("content") + .and_then(|c| c.get("name")) + .and_then(|v| v.as_str()) + .unwrap_or("unknown") + .to_string(); + Ok(ok_result(&format!("hello {name}"))) + }) + }); + + // Wait for the task to surface the input request. + let mut saw_input_required = false; + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + let detailed = manager.get_task(&task.task_id).unwrap(); + if let TaskPayload::InputRequired { input_requests } = &detailed.payload { + assert!(input_requests.contains_key("name-1")); + saw_input_required = true; + break; + } + } + assert!(saw_input_required, "task never reached input_required"); - /// Fetch the metadata for a running or recently completed task. - pub fn task_descriptor(&self, task_id: &str) -> Option<&OperationDescriptor> { - if let Some(task) = self.running_tasks.get(task_id) { - return Some(&task.descriptor); + // Respond via tasks/update. + manager + .update_task( + &task.task_id, + [( + "name-1".to_string(), + serde_json::json!({"action": "accept", "content": {"name": "Ada"}}), + )], + ) + .unwrap(); + + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + let detailed = manager.get_task(&task.task_id).unwrap(); + if detailed.status() == TaskStatus::Completed { + return; + } } - self.completed_results - .iter() - .rev() - .find(|result| result.descriptor.operation_id == task_id) - .map(|result| &result.descriptor) - } - - /// Attempt to cancel a running task. - pub fn cancel_task(&mut self, task_id: &str) -> bool { - self.collect_completed_results(); - if let Some(task) = self.running_tasks.remove(task_id) { - task.task_handle.abort(); - // Insert a cancelled result so callers can observe the terminal state. - let cancel_result = TaskResult { - descriptor: task.descriptor, - result: Err(Error::TaskError("Operation cancelled".to_string())), - }; - self.completed_results.push(cancel_result); - return true; + panic!("task did not complete after input response"); + } + + #[tokio::test] + async fn task_operation_reestablishes_request_association_scope() { + use crate::{ + model::RequestId, + service::{ORIGINATING_REQUEST, in_request_handler_scope}, + }; + + let manager = TaskManager::new(); + let observed = Arc::new(Mutex::new(None::)); + let observed_in_task = observed.clone(); + + ORIGINATING_REQUEST + .scope(RequestId::Number(7), async { + manager.spawn(TaskOptions::default(), move |_ctx| { + let observed_in_task = observed_in_task.clone(); + Box::pin(async move { + *observed_in_task.lock().unwrap() = Some(in_request_handler_scope()); + Ok(ok_result("done")) + }) + }) + }) + .await; + + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + if let Some(scoped) = *observed.lock().unwrap() { + assert!( + scoped, + "task operation must run inside the originating request's association scope" + ); + return; + } } - false - } - - /// Retrieve a completed task result if available. - pub fn take_completed_result(&mut self, task_id: &str) -> Option { - self.collect_completed_results(); - if let Some(position) = self - .completed_results - .iter() - .position(|result| result.descriptor.operation_id == task_id) - { - Some(self.completed_results.remove(position)) - } else { - None + panic!("task operation did not run"); + } + + #[tokio::test] + async fn task_operation_without_originating_request_is_unscoped() { + use crate::service::in_request_handler_scope; + + let manager = TaskManager::new(); + let observed = Arc::new(Mutex::new(None::)); + let observed_in_task = observed.clone(); + + manager.spawn(TaskOptions::default(), move |_ctx| { + let observed_in_task = observed_in_task.clone(); + Box::pin(async move { + *observed_in_task.lock().unwrap() = Some(in_request_handler_scope()); + Ok(ok_result("done")) + }) + }); + + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + if let Some(scoped) = *observed.lock().unwrap() { + assert!( + !scoped, + "task operation started without an originating request must remain unscoped" + ); + return; + } } + panic!("task operation did not run"); } } diff --git a/crates/rmcp/src/transport.rs b/crates/rmcp/src/transport.rs index 8cc48aa..fdedc9c 100644 --- a/crates/rmcp/src/transport.rs +++ b/crates/rmcp/src/transport.rs @@ -99,11 +99,12 @@ pub mod auth; pub use auth::JwtSigningAlgorithm; #[cfg(feature = "auth")] pub use auth::{ - AuthClient, AuthError, AuthorizationManager, AuthorizationSession, AuthorizedHttpClient, - ClientCredentialsConfig, CredentialStore, EXTENSION_OAUTH_CLIENT_CREDENTIALS, - InMemoryCredentialStore, InMemoryStateStore, OAuthHttpClient, OAuthHttpClientError, - OAuthHttpClientFuture, OAuthHttpRedirectPolicy, OAuthHttpRequest, ScopeUpgradeConfig, - StateStore, StoredAuthorizationState, StoredCredentials, WWWAuthenticateParams, + AuthClient, AuthError, AuthorizationManager, AuthorizationRequest, AuthorizationSession, + AuthorizedHttpClient, ClientCredentialsConfig, CredentialRefreshGuard, CredentialStore, + EXTENSION_OAUTH_CLIENT_CREDENTIALS, InMemoryCredentialStore, InMemoryStateStore, + OAuthHttpClient, OAuthHttpClientError, OAuthHttpClientFuture, OAuthHttpRedirectPolicy, + OAuthHttpRequest, ScopeUpgradeConfig, StateStore, StoredAuthorizationState, StoredCredentials, + WWWAuthenticateParams, }; // #[cfg(feature = "transport-ws")] @@ -271,6 +272,27 @@ impl DynamicTransportError { } } + pub(crate) fn is_authorization_required(&self) -> bool { + let mut error = Some(self.error.as_ref() as &(dyn std::error::Error + 'static)); + while let Some(current) = error { + #[cfg(feature = "auth")] + if matches!( + current.downcast_ref::(), + Some(auth::AuthError::AuthorizationRequired) + ) { + return true; + } + + #[cfg(feature = "transport-streamable-http-client")] + if current.is::() { + return true; + } + + error = current.source(); + } + false + } + pub fn downcast + 'static, R: ServiceRole>(self) -> Result { if !self.is::() { Err(self) diff --git a/crates/rmcp/src/transport/async_rw.rs b/crates/rmcp/src/transport/async_rw.rs index f50d913..e18fd4f 100644 --- a/crates/rmcp/src/transport/async_rw.rs +++ b/crates/rmcp/src/transport/async_rw.rs @@ -249,12 +249,13 @@ const UTF8_BOM: &[u8; 3] = b"\xEF\xBB\xBF"; /// Check if a method is a standard MCP method (request, response, or notification). /// This includes both requests and notifications defined in the MCP specification. /// -/// Based on MCP specification 2025-06-18: https://modelcontextprotocol.io/specification/2025-06-18 +/// Based on MCP specification 2026-07-28: https://modelcontextprotocol.io/specification/2026-07-28 fn is_standard_method(method: &str) -> bool { matches!( method, "initialize" | "ping" + | "server/discover" | "prompts/get" | "prompts/list" | "resources/list" @@ -262,9 +263,11 @@ fn is_standard_method(method: &str) -> bool { | "resources/subscribe" | "resources/unsubscribe" | "resources/templates/list" + | "subscriptions/listen" | "tools/call" | "tools/list" | "completion/complete" + | "elicitation/create" | "logging/setLevel" | "roots/list" | "sampling/createMessage" @@ -282,6 +285,7 @@ fn is_standard_notification(method: &str) -> bool { | "notifications/resources/list_changed" | "notifications/resources/updated" | "notifications/roots/list_changed" + | "notifications/subscriptions/acknowledged" | "notifications/tools/list_changed" ) } @@ -320,14 +324,12 @@ fn try_parse_with_compatibility( Ok(item) => Ok(Some(item)), Err(e) => { // Check if this is a notification that should be ignored for compatibility - if let Ok(json_value) = serde_json::from_str::(line_str) { - if let Some(method) = + if let Ok(json_value) = serde_json::from_str::(line_str) + && let Some(method) = json_value.get("method").and_then(serde_json::Value::as_str) - { - if should_ignore_notification(&json_value, method) { - return Ok(None); - } - } + && should_ignore_notification(&json_value, method) + { + return Ok(None); } tracing::debug!( @@ -584,6 +586,9 @@ mod test { assert!(is_standard_notification("notifications/tools/list_changed")); assert!(is_standard_notification("notifications/message")); assert!(is_standard_notification("notifications/roots/list_changed")); + assert!(is_standard_notification( + "notifications/subscriptions/acknowledged" + )); // Test that non-standard notifications are not recognized assert!(!is_standard_notification("notifications/stderr")); @@ -592,6 +597,39 @@ mod test { assert!(!is_standard_notification("some/other/method")); } + #[test] + fn test_2026_07_28_standard_method_check() { + for method in [ + "elicitation/create", + "server/discover", + "subscriptions/listen", + ] { + assert!(is_standard_method(method), "{method} should be standard"); + } + } + + #[test] + fn test_current_standard_notification_is_not_ignored() { + let method = "notifications/subscriptions/acknowledged"; + let notification = serde_json::json!({ + "jsonrpc": "2.0", + "method": method, + }); + + assert!(!should_ignore_notification(¬ification, method)); + } + + #[test] + fn test_unknown_notification_is_ignored() { + let method = "notifications/custom"; + let notification = serde_json::json!({ + "jsonrpc": "2.0", + "method": method, + }); + + assert!(should_ignore_notification(¬ification, method)); + } + #[test] fn test_compatibility_function() { // Test the compatibility function directly diff --git a/crates/rmcp/src/transport/auth.rs b/crates/rmcp/src/transport/auth.rs index 353c911..0727368 100644 --- a/crates/rmcp/src/transport/auth.rs +++ b/crates/rmcp/src/transport/auth.rs @@ -1,5 +1,5 @@ use std::{ - collections::HashMap, + collections::{HashMap, HashSet}, future::Future, net::{IpAddr, Ipv4Addr, Ipv6Addr}, pin::Pin, @@ -13,7 +13,8 @@ use oauth2::{ AsyncHttpClient, AuthType, AuthUrl, AuthorizationCode, ClientId, ClientSecret, CsrfToken, EmptyExtraTokenFields, ExtraTokenFields, HttpRequest, HttpResponse, PkceCodeChallenge, PkceCodeVerifier, RedirectUrl, RefreshToken, RequestTokenError, Scope, StandardTokenResponse, - TokenResponse, TokenUrl, basic::BasicTokenType, + TokenResponse, TokenUrl, + basic::{BasicErrorResponseType, BasicTokenType}, }; use reqwest::{ Client as ReqwestClient, IntoUrl, StatusCode, Url, @@ -27,6 +28,9 @@ use tracing::{debug, warn}; use crate::transport::common::http_header::HEADER_MCP_PROTOCOL_VERSION; +#[cfg(feature = "auth-enterprise-managed")] +pub mod enterprise; + const DEFAULT_HTTP_TIMEOUT: Duration = Duration::from_secs(30); const MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES: usize = 1024 * 1024; const MAX_OAUTH_DISCOVERY_REDIRECTS: usize = 10; @@ -69,20 +73,19 @@ impl OAuthHttpRequest { } } -/// Error returned by a custom OAuth HTTP client. -#[derive(Debug, Error)] -#[error("{message}")] -pub struct OAuthHttpClientError { - message: String, -} +/// Type-erased error returned by an [`OAuthHttpClient`]. +pub type OAuthHttpClientError = Box; -impl OAuthHttpClientError { - /// Create an error from a transport-provided message. - pub fn new(message: impl Into) -> Self { - Self { - message: message.into(), - } - } +#[derive(Debug, Error)] +enum OAuthHttpError { + #[error("OAuth HTTP response body exceeds {0} bytes")] + ResponseBodyTooLarge(usize), + #[error("unexpected HTTP status {0}")] + UnexpectedStatus(StatusCode), + #[error("OAuth discovery redirect to non-same-origin URL rejected: {0}")] + CrossOriginRedirect(Url), + #[error("OAuth discovery exceeded {0} redirects")] + TooManyRedirects(usize), } /// Future returned by [`OAuthHttpClient::execute`]. @@ -99,6 +102,19 @@ pub trait OAuthHttpClient: Send + Sync { fn execute(&self, request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_>; } +/// Create an OAuth HTTP client with the SDK's default reqwest configuration. +/// +/// Honors each request's redirect policy, with a 30-second timeout and bounded +/// response bodies. Enable a TLS feature such as `reqwest` for HTTPS requests. +/// Implement [`OAuthHttpClient`] instead when custom network policy is required. +pub fn default_oauth_http_client() -> Result { + let client = ReqwestClient::builder() + .timeout(DEFAULT_HTTP_TIMEOUT) + .build() + .map_err(|error| AuthError::InternalError(error.to_string()))?; + ReqwestOAuthHttpClient::new(client) +} + struct ReqwestOAuthHttpClient { follow_redirects: ReqwestClient, stop_redirects: ReqwestClient, @@ -131,11 +147,11 @@ impl OAuthHttpClient for ReqwestOAuthHttpClient { OAuthHttpRedirectPolicy::Stop => &self.stop_redirects, }; let request = reqwest::Request::try_from(request) - .map_err(|error| OAuthHttpClientError::new(error.to_string()))?; + .map_err(|error| Box::new(error) as OAuthHttpClientError)?; let response = client .execute(request) .await - .map_err(|error| OAuthHttpClientError::new(error.to_string()))?; + .map_err(|error| Box::new(error) as OAuthHttpClientError)?; let mut builder = oauth2::http::Response::builder() .status(response.status()) @@ -146,17 +162,17 @@ impl OAuthHttpClient for ReqwestOAuthHttpClient { let mut body = Vec::new(); let mut body_stream = response.bytes_stream(); while let Some(chunk) = body_stream.next().await { - let chunk = chunk.map_err(|error| OAuthHttpClientError::new(error.to_string()))?; + let chunk = chunk.map_err(|error| Box::new(error) as OAuthHttpClientError)?; if chunk.len() > MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES - body.len() { - return Err(OAuthHttpClientError::new(format!( - "OAuth HTTP response body exceeds {MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES} bytes" - ))); + return Err(Box::new(OAuthHttpError::ResponseBodyTooLarge( + MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES, + )) as OAuthHttpClientError); } body.extend_from_slice(&chunk); } builder .body(body) - .map_err(|error| OAuthHttpClientError::new(error.to_string())) + .map_err(|error| Box::new(error) as OAuthHttpClientError) }) } } @@ -166,16 +182,35 @@ struct OAuth2HttpClient<'a> { redirect_policy: OAuthHttpRedirectPolicy, } +#[derive(Debug)] +struct OAuth2HttpClientError(OAuthHttpClientError); + +impl std::fmt::Display for OAuth2HttpClientError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("OAuth HTTP request failed") + } +} + +impl std::error::Error for OAuth2HttpClientError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(self.0.as_ref()) + } +} + impl<'c> AsyncHttpClient<'c> for OAuth2HttpClient<'_> { - type Error = OAuthHttpClientError; + type Error = OAuth2HttpClientError; type Future = std::pin::Pin< Box> + Send + 'c>, >; fn call(&'c self, request: HttpRequest) -> Self::Future { - self.client - .execute(OAuthHttpRequest::new(request, self.redirect_policy)) + Box::pin(async move { + self.client + .execute(OAuthHttpRequest::new(request, self.redirect_policy)) + .await + .map_err(OAuth2HttpClientError) + }) } } @@ -194,6 +229,8 @@ pub struct StoredCredentials { pub granted_scopes: Vec, #[serde(default)] pub token_received_at: Option, + #[serde(default)] + pub issuer: Option, } impl std::fmt::Debug for StoredCredentials { @@ -206,6 +243,7 @@ impl std::fmt::Debug for StoredCredentials { ) .field("granted_scopes", &self.granted_scopes) .field("token_received_at", &self.token_received_at) + .field("issuer", &self.issuer) .finish() } } @@ -223,6 +261,30 @@ impl StoredCredentials { token_response, granted_scopes, token_received_at, + issuer: None, + } + } + + pub fn with_issuer(mut self, issuer: Option) -> Self { + self.issuer = issuer; + self + } +} + +/// An owned guard held across a credential refresh and its save. +/// +/// Stores can wrap a file lock, an owned mutex guard, or another coordination +/// primitive. Dropping this value releases the guard. +#[must_use = "dropping the guard releases refresh coordination"] +pub struct CredentialRefreshGuard { + _guard: Box, +} + +impl CredentialRefreshGuard { + /// Wrap a guard whose lifetime coordinates access to the stored credentials. + pub fn new(guard: impl Send + 'static) -> Self { + Self { + _guard: Box::new(guard), } } } @@ -232,6 +294,9 @@ impl StoredCredentials { /// Implementations of this trait can provide custom storage backends /// for OAuth2 credentials, such as file-based storage, keychain integration, /// or database storage. +/// +/// Return [`AuthError::CredentialStoreError`] for backend or locking failures +/// so they remain distinct from errors requiring reauthorization. #[async_trait] pub trait CredentialStore: Send + Sync { async fn load(&self) -> Result, AuthError>; @@ -239,6 +304,16 @@ pub trait CredentialStore: Send + Sync { async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError>; async fn clear(&self) -> Result<(), AuthError>; + + /// Optionally coordinate refreshes that share these credentials. + /// + /// The manager acquires this guard before loading credentials and retains it + /// through the token request and save. `load` and `save` must not reacquire + /// the same lock. Writers that bypass the guard are not coordinated with it. + /// The default does not coordinate refreshes. + async fn acquire_refresh_guard(&self) -> Result, AuthError> { + Ok(None) + } } /// In-memory credential store (default implementation) @@ -287,6 +362,10 @@ pub struct StoredAuthorizationState { #[serde(default)] pub require_issuer: bool, pub created_at: u64, + /// scopes requested in this round, used to resolve the grant when the token response omits + /// `scope` (RFC 6749 §5.1) + #[serde(default)] + pub requested_scopes: Vec, } impl std::fmt::Debug for StoredAuthorizationState { @@ -297,6 +376,7 @@ impl std::fmt::Debug for StoredAuthorizationState { .field("expected_issuer", &self.expected_issuer) .field("require_issuer", &self.require_issuer) .field("created_at", &self.created_at) + .field("requested_scopes", &self.requested_scopes) .finish() } } @@ -348,9 +428,16 @@ impl StoredAuthorizationState { .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_secs()) .unwrap_or(0), + requested_scopes: Vec::new(), } } + /// record the scopes requested in this authorization round (SEP-2350) + pub fn with_requested_scopes(mut self, scopes: Vec) -> Self { + self.requested_scopes = scopes; + self + } + pub fn into_pkce_verifier(self) -> PkceCodeVerifier { PkceCodeVerifier::new(self.pkce_verifier) } @@ -465,9 +552,19 @@ pub enum AuthError { #[error("OAuth token exchange failed: {0}")] TokenExchangeFailed(String), + /// The refresh attempt failed without a definitive refresh-token rejection. + /// + /// Callers may retry this error because it includes transient request and provider failures. #[error("OAuth token refresh failed: {0}")] TokenRefreshFailed(String), + /// The authorization server definitively rejected the refresh token. + #[error("OAuth refresh token was rejected: {0}")] + TokenRefreshRejected(String), + + #[error("OAuth credential store failed: {0}")] + CredentialStoreError(String), + #[error("HTTP error: {0}")] HttpError(#[from] reqwest::Error), @@ -477,6 +574,9 @@ pub enum AuthError { #[error("Metadata error: {0}")] MetadataError(String), + #[error("Authorization server does not support the required PKCE code challenge method (S256)")] + PkceUnsupported, + #[error("URL parse error: {0}")] UrlError(#[from] url::ParseError), @@ -540,6 +640,43 @@ pub struct AuthorizationMetadata { pub additional_fields: HashMap, } +/// How [`AuthorizationMetadata`] was obtained during discovery. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum AuthorizationMetadataSource { + /// Discovered through RFC 9728 protected resource metadata. + ProtectedResourceMetadata, + /// Discovered through RFC 8414 / OpenID Connect metadata at the server's + /// base URL. + AuthorizationServerMetadata, + /// Nothing was discovered; the endpoints were synthesized from the base + /// URL (`/authorize`, `/token`, `/register`) for compatibility with the + /// 2025-03-26 MCP spec's default-endpoint fallback. The server gave no + /// evidence that it supports OAuth. + /// + /// [Newer MCP revisions] require metadata discovery and do not define an + /// endpoint-synthesis fallback. + /// + /// [Newer MCP revisions]: https://modelcontextprotocol.io/specification/2026-07-28/basic/authorization/authorization-server-discovery#protected-resource-metadata-discovery-requirements + LegacyEndpointFallback, +} + +impl AuthorizationMetadataSource { + /// Whether the metadata was actually published by the server, as opposed + /// to synthesized by the client as a legacy compatibility fallback. + pub fn is_discovered(self) -> bool { + !matches!(self, Self::LegacyEndpointFallback) + } +} + +/// [`AuthorizationMetadata`] together with how it was resolved. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct AuthorizationMetadataResolution { + pub metadata: AuthorizationMetadata, + pub source: AuthorizationMetadataSource, +} + #[derive(Debug, Clone, Deserialize)] struct ResourceServerMetadata { resource: Option, @@ -548,6 +685,19 @@ struct ResourceServerMetadata { scopes_supported: Option>, } +/// How a url that may hold protected resource metadata was arrived at, which +/// decides what a document that is not metadata means there. +#[derive(Debug, Clone, Copy)] +enum ResourceMetadataUrlOrigin { + /// The resource itself named this url in the `resource_metadata` parameter + /// of its `WWW-Authenticate` challenge. + Advertised, + /// The url was derived from the base url on the chance that the document is + /// published there, or named by a challenge on such a url. Either way the + /// run only guessed its way here. + WellKnownGuess, +} + /// Parameters extracted from WWW-Authenticate header #[derive(Debug, Clone, Default)] #[non_exhaustive] @@ -559,6 +709,12 @@ pub struct WWWAuthenticateParams { } impl WWWAuthenticateParams { + /// Parse a `WWW-Authenticate` header value, resolving a relative + /// `resource_metadata` URL against `base_url`. + pub fn parse(header: &str, base_url: &Url) -> Self { + AuthorizationManager::extract_www_authenticate_params(header, base_url) + } + /// check if this is an insufficient_scope error pub fn is_insufficient_scope(&self) -> bool { self.error.as_deref() == Some("insufficient_scope") @@ -609,6 +765,142 @@ impl OAuthClientConfig { } } +/// Declarative description of the client identity material available for an +/// authorization flow. +/// +/// The [MCP authorization specification](https://modelcontextprotocol.io/specification/2026-07-28/basic/authorization/client-registration) +/// recommends that clients obtain a client ID using the following priority +/// order. [`OAuthState::start_authorization`] and [`AuthorizationSession::new`] +/// apply it internally: +/// +/// 1. Pre-registered client information +/// ([`with_preregistered_client`](Self::with_preregistered_client)), when +/// the client already holds a `client_id` issued out of band +/// 2. Client ID Metadata Documents (SEP-991, +/// [`with_client_metadata_url`](Self::with_client_metadata_url)), when the +/// authorization server advertises `client_id_metadata_document_supported` +/// 3. Dynamic Client Registration as a fallback, when the authorization server +/// advertises a `registration_endpoint` +/// +/// Provide whichever identity material the client has available; the SDK +/// selects the highest-priority mechanism the server supports. +/// +/// ```rust,ignore +/// let request = AuthorizationRequest::new("http://localhost:8080/callback") +/// // Omit `with_scopes` to let the SDK auto-select scopes from server metadata. +/// // Used when the server supports CIMD and no pre-registered client is set. +/// .with_client_metadata_url("https://example.com/client-metadata.json") +/// // used for dynamic client registration as a last resort +/// .with_client_name("My MCP Client"); +/// oauth_state.start_authorization(request).await?; +/// ``` +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct AuthorizationRequest { + /// Redirect URI for the authorization-code flow. + pub redirect_uri: String, + /// Scopes to request. When empty, the SDK selects scopes from the + /// server's `WWW-Authenticate` challenge, Protected Resource Metadata, + /// or authorization server metadata. + pub scopes: Vec, + /// Human-readable client name, used for Dynamic Client Registration. + pub client_name: Option, + /// Pre-registered client ID obtained from the authorization server out of + /// band. When set, registration is skipped entirely. + pub client_id: Option, + /// Client secret paired with the pre-registered [`client_id`](Self::client_id). + pub client_secret: Option, + /// HTTPS URL of a Client ID Metadata Document (SEP-991). Used when the + /// authorization server advertises `client_id_metadata_document_supported` + /// and no pre-registered client is configured. + pub client_metadata_url: Option, + /// OIDC Dynamic Client Registration `application_type` (SEP-837), + /// e.g. `"native"` or `"web"`. + pub application_type: Option, + /// `WWW-Authenticate` header value from a real request's 401 response. + /// When set, discovery is seeded from the challenge (its + /// `resource_metadata` URL and `scope` hint) instead of probing the + /// server — the reactive discovery path. + pub challenge: Option, +} + +impl AuthorizationRequest { + /// Create a request for the given redirect URI. With no further identity + /// material, authorization falls back to Dynamic Client Registration. + pub fn new(redirect_uri: impl Into) -> Self { + Self { + redirect_uri: redirect_uri.into(), + scopes: Vec::new(), + client_name: None, + client_id: None, + client_secret: None, + client_metadata_url: None, + application_type: None, + challenge: None, + } + } + + /// Set the scopes to request. When not set, the SDK auto-selects scopes + /// using its normal scope-selection policy. + pub fn with_scopes(mut self, scopes: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.scopes = scopes.into_iter().map(Into::into).collect(); + self + } + + /// Set the client name used for Dynamic Client Registration. + pub fn with_client_name(mut self, client_name: impl Into) -> Self { + self.client_name = Some(client_name.into()); + self + } + + /// Use a client ID that was pre-registered with the authorization server + /// out of band. This takes priority over every other mechanism. + /// + /// Pair with [`with_client_secret`](Self::with_client_secret) for + /// confidential clients. + pub fn with_preregistered_client(mut self, client_id: impl Into) -> Self { + self.client_id = Some(client_id.into()); + self + } + + /// Set the client secret paired with a pre-registered client ID. + /// + /// Must be used together with + /// [`with_preregistered_client`](Self::with_preregistered_client); + /// authorization fails with [`AuthError::RegistrationFailed`] if a secret + /// is provided without a client ID. + pub fn with_client_secret(mut self, client_secret: impl Into) -> Self { + self.client_secret = Some(client_secret.into()); + self + } + + /// Set the HTTPS URL of a Client ID Metadata Document (SEP-991). Used when + /// the authorization server advertises support and no pre-registered + /// client is configured. + pub fn with_client_metadata_url(mut self, client_metadata_url: impl Into) -> Self { + self.client_metadata_url = Some(client_metadata_url.into()); + self + } + + /// Set the OIDC Dynamic Client Registration `application_type` (SEP-837), + /// e.g. `"native"` or `"web"`. + pub fn with_application_type(mut self, application_type: impl Into) -> Self { + self.application_type = Some(application_type.into()); + self + } + + /// Seed discovery from the `WWW-Authenticate` header value of a real + /// request's 401 response instead of probing the server. + pub fn with_challenge(mut self, www_authenticate: impl Into) -> Self { + self.challenge = Some(www_authenticate.into()); + self + } +} + // add type aliases for oauth2 types type OAuthErrorResponse = oauth2::StandardErrorResponse; @@ -685,7 +977,7 @@ impl JwtSigningAlgorithm { /// This supports two authentication methods: /// - `ClientSecret`: credentials sent in the request body /// - `PrivateKeyJwt`: RFC 7523 signed JWT assertion (requires `auth-client-credentials-jwt` feature) -#[derive(Debug, Clone)] +#[derive(Clone)] #[non_exhaustive] pub enum ClientCredentialsConfig { /// Client secret authentication (credentials in request body) @@ -701,13 +993,71 @@ pub enum ClientCredentialsConfig { client_id: String, signing_key: Vec, signing_algorithm: JwtSigningAlgorithm, - /// Override the `aud` claim in the JWT assertion; defaults to token_endpoint + /// Overrides the authorization server issuer used for the JWT `aud` claim. token_endpoint_audience: Option, scopes: Vec, resource: Option, }, } +impl std::fmt::Debug for ClientCredentialsConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::ClientSecret { + client_id, + scopes, + resource, + .. + } => f + .debug_struct("ClientSecret") + .field("client_id", client_id) + .field("client_secret", &"") + .field("scopes", scopes) + .field("resource", resource) + .finish(), + #[cfg(feature = "auth-client-credentials-jwt")] + Self::PrivateKeyJwt { + client_id, + signing_algorithm, + token_endpoint_audience, + scopes, + resource, + .. + } => f + .debug_struct("PrivateKeyJwt") + .field("client_id", client_id) + .field("signing_key", &"") + .field("signing_algorithm", signing_algorithm) + .field("token_endpoint_audience", token_endpoint_audience) + .field("scopes", scopes) + .field("resource", resource) + .finish(), + } + } +} + +#[cfg(feature = "auth-client-credentials-jwt")] +fn client_authentication_audience<'a>( + metadata: &'a AuthorizationMetadata, + configured_audience: Option<&'a str>, +) -> Result<&'a str, AuthError> { + configured_audience + .or(metadata.issuer.as_deref()) + .ok_or_else(|| { + AuthError::ClientCredentialsError( + "Authorization server metadata is missing the issuer required for private_key_jwt" + .to_string(), + ) + }) +} + +#[cfg(feature = "auth-client-credentials-jwt")] +fn client_authentication_header(algorithm: JwtSigningAlgorithm) -> jsonwebtoken::Header { + let mut header = jsonwebtoken::Header::new(algorithm.to_jsonwebtoken_algorithm()); + header.typ = Some("client-authentication+jwt".to_string()); + header +} + impl ClientCredentialsConfig { fn client_id(&self) -> &str { match self { @@ -778,8 +1128,11 @@ pub struct AuthorizationManager { www_auth_scopes: RwLock>, /// scopes_supported from protected resource metadata (RFC 9728) resource_scopes: RwLock>, + /// resource indicator from protected resource metadata, used for RFC 8707 `resource` + discovered_resource: RwLock>, /// OIDC Dynamic Client Registration `application_type` (SEP-837) application_type: Option, + allow_missing_issuer: bool, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -828,6 +1181,18 @@ fn is_https_url(value: &str) -> bool { .unwrap_or(false) } +#[cfg(feature = "auth-client-credentials-jwt")] +fn is_allowed_client_credentials_endpoint(resource: &Url, token_endpoint: &Url) -> bool { + token_endpoint.scheme() == "https" + || (token_endpoint.scheme() == "http" + && resource + .host_str() + .is_some_and(AuthorizationManager::is_loopback_metadata_host) + && token_endpoint + .host_str() + .is_some_and(AuthorizationManager::is_loopback_metadata_host)) +} + impl AuthorizationManager { fn is_http_url(url: &Url) -> bool { matches!(url.scheme(), "http" | "https") && url.host_str().is_some() @@ -891,11 +1256,31 @@ impl AuthorizationManager { } } - fn is_allowed_authorization_server_metadata_url(url: &Url) -> bool { - Self::is_http_url(url) - && url - .host_str() - .is_some_and(|host| !Self::is_disallowed_metadata_host(host)) + fn is_loopback_metadata_host(host: &str) -> bool { + let host = host.trim_end_matches('.').to_ascii_lowercase(); + host == "localhost" + || host.ends_with(".localhost") + || matches!(host.parse::(), Ok(IpAddr::V4(addr)) if addr.is_loopback()) + || matches!(host.parse::(), Ok(IpAddr::V6(addr)) if addr.is_loopback()) + } + + fn is_allowed_authorization_server_metadata_url(base_url: &Url, url: &Url) -> bool { + if !Self::is_http_url(url) { + return false; + } + + let Some(host) = url.host_str() else { + return false; + }; + + if !Self::is_disallowed_metadata_host(host) { + return true; + } + + base_url + .host_str() + .is_some_and(Self::is_loopback_metadata_host) + && Self::is_loopback_metadata_host(host) } fn resolve_resource_metadata_url(value: &str, base_url: &Url) -> Option { @@ -952,13 +1337,9 @@ impl AuthorizationManager { /// create new auth manager with base url pub async fn new(base_url: U) -> Result { - let http_client = ReqwestClient::builder() - .timeout(DEFAULT_HTTP_TIMEOUT) - .build() - .map_err(|e| AuthError::InternalError(e.to_string()))?; Self::new_inner( base_url, - Arc::new(ReqwestOAuthHttpClient::new(http_client)?), + Arc::new(default_oauth_http_client()?), OAuthHttpRedirectPolicy::Stop, ) .await @@ -992,7 +1373,9 @@ impl AuthorizationManager { scope_upgrade_config: ScopeUpgradeConfig::default(), www_auth_scopes: RwLock::new(Vec::new()), resource_scopes: RwLock::new(Vec::new()), + discovered_resource: RwLock::new(None), application_type: Some(DEFAULT_APPLICATION_TYPE.to_string()), + allow_missing_issuer: false, }; Ok(manager) @@ -1003,6 +1386,17 @@ impl AuthorizationManager { self.scope_upgrade_config = config; } + /// Configure whether authorization server metadata discovery tolerates a + /// missing `issuer` field. + /// + /// The default is `false`, enforcing the RFC 8414/OIDC requirement that + /// discovered metadata include `issuer` whenever the expected issuer can be + /// derived from the discovery URL. Set this to `true` only for compatibility + /// with authorization servers that return incomplete metadata. + pub fn set_allow_missing_issuer(&mut self, allow: bool) { + self.allow_missing_issuer = allow; + } + /// Set a custom credential store /// /// This allows you to provide your own implementation of credential storage, @@ -1021,8 +1415,10 @@ impl AuthorizationManager { /// Set OAuth2 authorization metadata /// - /// This should be called after discovering metadata via `discover_metadata()` - /// and before creating an `AuthorizationSession`. + /// This should be called with + /// [`AuthorizationMetadataResolution::metadata`] after + /// [`Self::resolve_metadata`] and before creating an + /// [`AuthorizationSession`]. pub fn set_metadata(&mut self, metadata: AuthorizationMetadata) { self.metadata = Some(metadata); } @@ -1030,18 +1426,59 @@ impl AuthorizationManager { /// Initialize from stored credentials if available /// /// This will load credentials from the credential store and configure - /// the client if credentials are found. + /// the client if credentials are found. Returns `false` when credentials + /// are absent or discarded after an authorization-server change. pub async fn initialize_from_store(&mut self) -> Result { - if let Some(stored) = self.credential_store.load().await? { - if stored.token_response.is_some() { - if self.metadata.is_none() { - let metadata = self.discover_metadata().await?; - self.metadata = Some(metadata); - } - - self.configure_client_id(&stored.client_id)?; - return Ok(true); + if let Some(stored) = self.credential_store.load().await? + && stored.token_response.is_some() + { + if self.metadata.is_none() { + let resolution = self.resolve_metadata().await?; + self.metadata = Some(resolution.metadata); } + + if let (Some(stored_issuer), Some(current_issuer)) = + (stored.issuer.as_deref(), self.metadata_issuer().as_deref()) + { + // A CIMD client ID is the client's metadata URL, so it is + // portable across authorization servers and exempt here. + if stored_issuer != current_issuer { + if is_https_url(&stored.client_id) { + // A CIMD client ID is the client's metadata URL, so it is + // portable across authorization servers — but the tokens + // were minted by the previous AS and must not be reused. + tracing::warn!( + stored_issuer, + current_issuer, + "authorization server issuer changed; discarding tokens but keeping portable CIMD client ID" + ); + self.credential_store + .save( + StoredCredentials::new( + stored.client_id.clone(), + None, + vec![], + None, + ) + .with_issuer(self.metadata_issuer()), + ) + .await?; + self.configure_client_id(&stored.client_id)?; + return Ok(false); + } + + tracing::warn!( + stored_issuer, + current_issuer, + "authorization server issuer changed; clearing stored credentials bound to the previous issuer" + ); + self.credential_store.clear().await?; + return Ok(false); + } + } + + self.configure_client_id(&stored.client_id)?; + return Ok(true); } Ok(false) } @@ -1064,19 +1501,114 @@ impl AuthorizationManager { Ok(()) } - /// discover oauth2 metadata (per SEP-985: Protected Resource Metadata first, then direct OAuth) - pub async fn discover_metadata(&self) -> Result { + /// Resolve OAuth 2.0 metadata and report how it was obtained. + /// + /// Discovery follows SEP-985: protected resource metadata first, then + /// direct OAuth 2.0 Authorization Server Metadata or OpenID Connect + /// Discovery. When discovery finds nothing, the result contains legacy + /// default endpoints derived from the base URL and + /// [`AuthorizationMetadataSource::LegacyEndpointFallback`]. + /// + /// # Examples + /// + /// ```no_run + /// use rmcp::transport::auth::{AuthorizationManager, AuthorizationMetadataSource}; + /// + /// # async fn resolve() -> Result<(), Box> { + /// let mut manager = AuthorizationManager::new("https://mcp.example.com").await?; + /// let resolution = manager.resolve_metadata().await?; + /// + /// if resolution.source == AuthorizationMetadataSource::LegacyEndpointFallback { + /// println!("the server did not publish OAuth metadata"); + /// } + /// + /// manager.set_metadata(resolution.metadata); + /// # Ok(()) + /// # } + /// ``` + pub async fn resolve_metadata(&self) -> Result { if let Some(metadata) = self.discover_oauth_server_via_resource_metadata().await? { - return Ok(metadata); + return Ok(AuthorizationMetadataResolution { + metadata, + source: AuthorizationMetadataSource::ProtectedResourceMetadata, + }); } - if let Some(metadata) = self.try_discover_oauth_server(&self.base_url).await? { - return Ok(metadata); + if let Some(metadata) = self.try_discover_oauth_server(&self.base_url, None).await? { + return Ok(AuthorizationMetadataResolution { + metadata, + source: AuthorizationMetadataSource::AuthorizationServerMetadata, + }); } - // No valid authorization metadata found - return error instead of guessing - // OAuth endpoints must be discovered from the server, not constructed by the client - Err(AuthError::NoAuthorizationSupport) + debug!("falling back to legacy OAuth endpoints derived from the base URL"); + Ok(AuthorizationMetadataResolution { + metadata: Self::legacy_authorization_metadata(&self.base_url), + source: AuthorizationMetadataSource::LegacyEndpointFallback, + }) + } + + /// Resolve authorization server metadata starting from the + /// `WWW-Authenticate` challenge of a real request's 401 response — the + /// reactive discovery path, matching the TypeScript and Python SDKs. + /// + /// Seeds scope selection with the challenge's `scope` hint and prefers + /// the challenge's `resource_metadata` URL; falls back to + /// [`resolve_metadata`](Self::resolve_metadata) when the challenge is + /// `None` or carries no usable metadata pointer. + pub async fn resolve_metadata_from_challenge( + &self, + www_authenticate: Option<&str>, + ) -> Result { + let Some(www_authenticate) = www_authenticate else { + return self.resolve_metadata().await; + }; + let params = WWWAuthenticateParams::parse(www_authenticate, &self.base_url); + + self.record_challenge_scope(¶ms).await; + + if let Some(resource_metadata_url) = ¶ms.resource_metadata_url + && let Some(metadata) = self + .discover_oauth_server_from_resource_metadata_url(resource_metadata_url) + .await? + { + return Ok(AuthorizationMetadataResolution { + metadata, + source: AuthorizationMetadataSource::ProtectedResourceMetadata, + }); + } + + self.resolve_metadata().await + } + + /// Store a challenge's `scope` hint for later scope selection. + async fn record_challenge_scope(&self, params: &WWWAuthenticateParams) { + let Some(scope) = ¶ms.scope else { + return; + }; + let scopes: Vec = scope.split_whitespace().map(str::to_string).collect(); + if scopes.is_empty() { + return; + } + debug!("WWW-Authenticate challenge contains scope: {scope}"); + *self.www_auth_scopes.write().await = scopes; + } + + fn legacy_authorization_metadata(base_url: &Url) -> AuthorizationMetadata { + let endpoint = |path: &str| { + let mut url = base_url.clone(); + url.set_query(None); + url.set_fragment(None); + url.set_path(path); + url.to_string() + }; + + AuthorizationMetadata { + authorization_endpoint: endpoint("/authorize"), + token_endpoint: endpoint("/token"), + registration_endpoint: Some(endpoint("/register")), + ..Default::default() + } } /// get client id and credentials @@ -1153,24 +1685,27 @@ impl AuthorizationManager { // RFC 8414 RECOMMENDS response_types_supported in the metadata. This field is optional, // but if present and does not include the flow we use ("code"), bail out early with a clear error. - if let Some(response_types_supported) = metadata.response_types_supported.as_ref() { - if !response_types_supported.contains(&response_type.to_string()) { - return Err(AuthError::InvalidScope(response_type.to_string())); - } + if let Some(response_types_supported) = metadata.response_types_supported.as_ref() + && !response_types_supported.contains(&response_type.to_string()) + { + return Err(AuthError::InvalidScope(response_type.to_string())); } - // for PKCE, we always send s256 since oauth 2.1 requires servers to support it, - // but warn if the server metadata suggests otherwise + // The client always sends an S256 challenge. A server that advertises + // methods without S256 can't do the flow we require, so refuse it. A + // server that omits the field is tolerated: it usually means the server + // didn't advertise PKCE, not that it lacks S256. match &metadata.code_challenge_methods_supported { Some(methods) if !methods.iter().any(|m| m == "S256") => { + return Err(AuthError::PkceUnsupported); + } + None => { warn!( - ?methods, - "server does not advertise S256 in code_challenge_methods_supported, \ - proceeding with S256 anyway as oauth 2.1 requires it. \ - The server is not compliant with the specification!" + "authorization server metadata omits code_challenge_methods_supported; \ + proceeding with an S256 challenge anyway" ); } - _ => {} + Some(_) => {} } Ok(()) @@ -1308,7 +1843,7 @@ impl AuthorizationManager { let mut auth_request = oauth_client .authorize_url(CsrfToken::new_random) .set_pkce_challenge(pkce_challenge) - .add_extra_param("resource", self.base_url.to_string()); + .add_extra_param("resource", self.oauth_resource().await); // add request scopes for scope in scopes { @@ -1317,7 +1852,7 @@ impl AuthorizationManager { let (auth_url, csrf_token) = auth_request.url(); - // store pkce verifier and expected issuer for later use via state store + // store pkce verifier, expected issuer, and the requested scopes for later use via state store let expected_issuer = self .metadata .as_ref() @@ -1337,7 +1872,8 @@ impl AuthorizationManager { &csrf_token, expected_issuer, require_issuer, - ); + ) + .with_requested_scopes(scopes.iter().map(|s| s.to_string()).collect()); self.state_store .save(csrf_token.secret(), stored_state) .await?; @@ -1345,6 +1881,14 @@ impl AuthorizationManager { Ok(auth_url.to_string()) } + async fn oauth_resource(&self) -> String { + self.discovered_resource + .read() + .await + .clone() + .unwrap_or_else(|| self.base_url.to_string()) + } + /// get the current granted scopes pub async fn get_current_scopes(&self) -> Vec { self.current_scopes.read().await.clone() @@ -1352,11 +1896,33 @@ impl AuthorizationManager { /// compute the union of current scopes and required scopes fn compute_scope_union(current: &[String], required: &str) -> Vec { - let mut scope_set: std::collections::HashSet = current.iter().cloned().collect(); - for scope in required.split_whitespace() { - scope_set.insert(scope.to_string()); + let mut scopes = current.to_vec(); + scopes.extend(required.split_whitespace().map(|s| s.to_string())); + Self::dedup_scopes(scopes) + } + + /// deduplicate scopes preserving first-seen order (SEP-2350: stable for testability) + fn dedup_scopes(scopes: Vec) -> Vec { + let mut seen = std::collections::HashSet::new(); + scopes + .into_iter() + .filter(|s| seen.insert(s.clone())) + .collect() + } + + /// resolve the granted scope set from a token response (SEP-2350, RFC 6749 §5.1): an explicit + /// `scope` may narrow the grant; an omitted one means the request was granted in full, so fall + /// back to the requested scopes (or the previously granted set when none were recorded). + fn resolve_granted_scopes( + response_scopes: Option>, + requested_scopes: &[String], + current_scopes: &[String], + ) -> Vec { + match response_scopes { + Some(scopes) => scopes, + None if !requested_scopes.is_empty() => requested_scopes.to_vec(), + None => current_scopes.to_vec(), } - scope_set.into_iter().collect() } /// check if a scope upgrade is possible and allowed @@ -1379,41 +1945,47 @@ impl AuthorizationManager { scopes } - /// select scopes based on SEP-835 priority: - /// 1. scope from WWW-Authenticate header (argument or stored from initial 401 probe) - /// 2. scopes_supported from protected resource metadata (RFC 9728) - /// 3. scopes_supported from authorization server metadata - /// 4. provided default scopes + /// select scopes following SEP-2350: re-authorization requests the union of the + /// previously requested scopes and the newly challenged scopes. Server-reported + /// scopes (WWW-Authenticate challenge, protected resource metadata) are operational + /// requirements for the current operation, never an exclusive directive, so they + /// accumulate rather than replace. The AS metadata and caller defaults only seed the + /// request when nothing has been requested or challenged yet. fn select_base_scopes( &self, www_authenticate_scope: Option<&str>, default_scopes: &[&str], ) -> Vec { + let mut accumulated: Vec = Vec::new(); + + // previously requested scopes + if let Ok(guard) = self.current_scopes.try_read() { + accumulated.extend(guard.iter().cloned()); + } + + // newly challenged scopes for the current operation (RFC 6750 §3.1) if let Some(scope) = www_authenticate_scope { - return scope.split_whitespace().map(|s| s.to_string()).collect(); + accumulated.extend(scope.split_whitespace().map(|s| s.to_string())); } - - // use scopes from initial 401 WWW-Authenticate header if let Ok(guard) = self.www_auth_scopes.try_read() { - if !guard.is_empty() { - return guard.clone(); - } + accumulated.extend(guard.iter().cloned()); } - // use scopes_supported from protected resource metadata (RFC 9728) + // scopes required for the current operation per protected resource metadata (RFC 9728) if let Ok(guard) = self.resource_scopes.try_read() { - if !guard.is_empty() { - return guard.clone(); - } + accumulated.extend(guard.iter().cloned()); } - // use scopes_supported from authorization server metadata - if let Some(metadata) = &self.metadata { - if let Some(scopes_supported) = &metadata.scopes_supported { - if !scopes_supported.is_empty() { - return scopes_supported.clone(); - } - } + if !accumulated.is_empty() { + return Self::dedup_scopes(accumulated); + } + + // nothing requested or challenged yet: seed from AS metadata, then caller defaults + if let Some(metadata) = &self.metadata + && let Some(scopes_supported) = &metadata.scopes_supported + && !scopes_supported.is_empty() + { + return scopes_supported.clone(); } default_scopes.iter().map(|s| s.to_string()).collect() @@ -1425,12 +1997,11 @@ impl AuthorizationManager { if scopes.is_empty() || scopes.iter().any(|s| s == "offline_access") { return; } - if let Some(metadata) = &self.metadata { - if let Some(supported) = &metadata.scopes_supported { - if supported.iter().any(|s| s == "offline_access") { - scopes.push("offline_access".to_string()); - } - } + if let Some(metadata) = &self.metadata + && let Some(supported) = &metadata.scopes_supported + && supported.iter().any(|s| s == "offline_access") + { + scopes.push("offline_access".to_string()); } } @@ -1542,10 +2113,15 @@ impl AuthorizationManager { AuthError::InternalError("Authorization state not found".to_string()) })?; - // Delete state after retrieval (one-time use) + Self::validate_authorization_response_issuer(&stored_state, received_issuer)?; + + // Consume state only after the callback is bound to the expected issuer. + // A callback with the correct state but a forged or missing required `iss` + // must not discard the PKCE verifier needed by the legitimate callback. self.state_store.delete(csrf_token).await?; - Self::validate_authorization_response_issuer(&stored_state, received_issuer)?; + // capture requested scopes before the state is consumed + let requested_scopes = stored_state.requested_scopes.clone(); // Reconstruct the PKCE verifier let pkce_verifier = stored_state.into_pkce_verifier(); @@ -1556,7 +2132,7 @@ impl AuthorizationManager { let token_result = match oauth_client .exchange_code(AuthorizationCode::new(code.to_string())) .set_pkce_verifier(pkce_verifier) - .add_extra_param("resource", self.base_url.to_string()) + .add_extra_param("resource", self.oauth_resource().await) .request_async(&OAuth2HttpClient { client: self.http_client.as_ref(), redirect_policy: OAuthHttpRedirectPolicy::Stop, @@ -1584,10 +2160,14 @@ impl AuthorizationManager { debug!("exchange token result: {:?}", token_result); - let granted_scopes: Vec = token_result + // SEP-2350: an omitted `scope` means the grant equals the request (RFC 6749 §5.1). + let response_scopes = token_result .scopes() - .map(|scopes| scopes.iter().map(|s| s.to_string()).collect()) - .unwrap_or_default(); + .map(|scopes| scopes.iter().map(|s| s.to_string()).collect()); + let granted_scopes = { + let current = self.current_scopes.read().await; + Self::resolve_granted_scopes(response_scopes, &requested_scopes, ¤t) + }; *self.current_scopes.write().await = granted_scopes.clone(); *self.scope_upgrade_attempts.write().await = 0; @@ -1598,6 +2178,7 @@ impl AuthorizationManager { token_response: Some(token_result.clone()), granted_scopes, token_received_at: Some(Self::now_epoch_secs()), + issuer: self.metadata_issuer(), }; self.credential_store.save(stored).await?; @@ -1611,15 +2192,21 @@ impl AuthorizationManager { .as_secs() } + fn metadata_issuer(&self) -> Option { + self.metadata.as_ref().and_then(|m| m.issuer.clone()) + } + /// Proactive refresh buffer: refresh tokens this many seconds before they expire /// to avoid races between token retrieval and the actual HTTP request. const REFRESH_BUFFER_SECS: u64 = 30; /// Get access token from local credential store. /// If expired, refresh it automatically when a refresh token is available. - /// When the access token has expired and no refresh token is available (or - /// the refresh itself fails), returns [`AuthError::AuthorizationRequired`] - /// so the caller can re-authenticate. + /// When the access token has expired and no refresh token is available, or the + /// authorization server rejects the refresh token, returns + /// [`AuthError::AuthorizationRequired`] so the caller can re-authenticate. + /// Transient refresh failures return [`AuthError::TokenRefreshFailed`] so the + /// caller can retry; other errors are propagated as-is. pub async fn get_access_token(&self) -> Result { let stored = self.credential_store.load().await?; let Some(stored_creds) = stored else { @@ -1654,13 +2241,13 @@ impl AuthorizationManager { /// refresh token or the server rejected it, return `AuthorizationRequired` /// so the caller can re-prompt the user. Infrastructure errors (e.g. store /// I/O failures, misconfigured client) are propagated as-is. - async fn try_refresh_or_reauth(&self) -> Result { + pub(crate) async fn try_refresh_or_reauth(&self) -> Result { match self.refresh_token().await { Ok(new_creds) => { tracing::info!("Refreshed access token."); Ok(new_creds.access_token().secret().to_string()) } - Err(e @ (AuthError::AuthorizationRequired | AuthError::TokenRefreshFailed(_))) => { + Err(e @ (AuthError::AuthorizationRequired | AuthError::TokenRefreshRejected(_))) => { tracing::warn!(error = %e, "Token refresh not possible, re-authorization required."); Err(AuthError::AuthorizationRequired) } @@ -1675,21 +2262,38 @@ impl AuthorizationManager { .as_ref() .ok_or_else(|| AuthError::InternalError("OAuth client not configured".to_string()))?; + // Held for the rest of this function so the load, the exchange, and the + // save stay inside one guarded section. + let _refresh_guard = self.credential_store.acquire_refresh_guard().await?; let stored = self.credential_store.load().await?; let stored_credentials = stored.ok_or(AuthError::AuthorizationRequired)?; + // Refreshing with another client's stored token would put that token on a + // request authenticated as this client. + if stored_credentials.client_id != oauth_client.client_id().as_str() { + tracing::warn!( + stored_client_id = stored_credentials.client_id.as_str(), + configured_client_id = oauth_client.client_id().as_str(), + "stored credentials belong to a different client; reauthorization required" + ); + return Err(AuthError::AuthorizationRequired); + } let current_credentials = stored_credentials .token_response .ok_or(AuthError::AuthorizationRequired)?; - let refresh_token = current_credentials.refresh_token().ok_or_else(|| { - AuthError::TokenRefreshFailed("No refresh token available".to_string()) - })?; + let refresh_token = current_credentials + .refresh_token() + .ok_or(AuthError::AuthorizationRequired)?; debug!("refresh token present, attempting refresh"); let refresh_token_value = RefreshToken::new(refresh_token.secret().to_string()); - let mut refresh_request = oauth_client.exchange_refresh_token(&refresh_token_value); + let mut refresh_request = oauth_client + .exchange_refresh_token(&refresh_token_value) + // RFC 8707: the resource indicator is required on token requests, including refreshes + .add_extra_param("resource", self.oauth_resource().await); let mut refresh_scopes = stored_credentials.granted_scopes; self.add_offline_access_if_supported(&mut refresh_scopes); + let requested_scopes = refresh_scopes.clone(); for scope in refresh_scopes { refresh_request = refresh_request.add_scope(Scope::new(scope)); } @@ -1699,7 +2303,14 @@ impl AuthorizationManager { redirect_policy: self.refresh_redirect_policy, }) .await - .map_err(|e| AuthError::TokenRefreshFailed(e.to_string()))?; + .map_err(|error| match &error { + RequestTokenError::ServerResponse(response) + if response.error() == &BasicErrorResponseType::InvalidGrant => + { + AuthError::TokenRefreshRejected(error.to_string()) + } + _ => AuthError::TokenRefreshFailed(error.to_string()), + })?; // RFC 6749 section 6: issuing a new refresh token on refresh is optional. // When the response omits one, keep the existing refresh token rather than @@ -1708,9 +2319,12 @@ impl AuthorizationManager { token_result.set_refresh_token(Some(refresh_token_value)); } - let granted_scopes: Vec = match token_result.scopes() { - Some(scopes) => scopes.iter().map(|s| s.to_string()).collect(), - None => self.current_scopes.read().await.clone(), + let response_scopes = token_result + .scopes() + .map(|scopes| scopes.iter().map(|s| s.to_string()).collect()); + let granted_scopes = { + let current = self.current_scopes.read().await; + Self::resolve_granted_scopes(response_scopes, &requested_scopes, ¤t) }; *self.current_scopes.write().await = granted_scopes.clone(); @@ -1721,6 +2335,7 @@ impl AuthorizationManager { token_response: Some(token_result.clone()), granted_scopes, token_received_at: Some(Self::now_epoch_secs()), + issuer: self.metadata_issuer(), }; self.credential_store.save(stored).await?; @@ -1797,9 +2412,13 @@ impl AuthorizationManager { async fn try_discover_oauth_server( &self, base_url: &Url, + expected_issuer: Option<&str>, ) -> Result, AuthError> { for discovery_url in Self::generate_discovery_urls(base_url) { - if let Some(metadata) = self.fetch_authorization_metadata(&discovery_url).await? { + if let Some(metadata) = self + .fetch_authorization_metadata(&discovery_url, expected_issuer) + .await? + { return Ok(Some(metadata)); } } @@ -1809,15 +2428,13 @@ impl AuthorizationManager { async fn fetch_authorization_metadata( &self, discovery_url: &Url, + expected_issuer: Option<&str>, ) -> Result, AuthError> { debug!("discovery url: {:?}", discovery_url); - let response = match self.discovery_get(discovery_url).await { - Ok(r) => r, - Err(e) => { - debug!("discovery request failed: {}", e); - return Ok(None); - } - }; + let response = self + .discovery_get(discovery_url) + .await + .map_err(|error| Self::discovery_failed(discovery_url, error))?; if response.status() != StatusCode::OK { debug!("discovery returned non-200: {}", response.status()); @@ -1825,7 +2442,14 @@ impl AuthorizationManager { } match serde_json::from_slice::(response.body()) { - Ok(metadata) => Ok(Some(metadata)), + Ok(metadata) => { + self.validate_authorization_metadata_issuer( + discovery_url, + expected_issuer, + &metadata, + )?; + Ok(Some(metadata)) + } Err(err) => { debug!("Failed to parse metadata for {}: {}", discovery_url, err); Ok(None) // malformed JSON ⇒ try next candidate @@ -1833,44 +2457,169 @@ impl AuthorizationManager { } } + fn expected_issuer_for_authorization_metadata_url(discovery_url: &Url) -> Option { + let path = discovery_url.path(); + let oauth_prefix = "/.well-known/oauth-authorization-server"; + let oidc_prefix = "/.well-known/openid-configuration"; + + let issuer_path = if path == oauth_prefix || path == oidc_prefix { + "" + } else if let Some(suffix) = path.strip_prefix(&format!("{oauth_prefix}/")) { + // RFC 8414 path-insertion form + suffix + } else if let Some(suffix) = path.strip_prefix(&format!("{oidc_prefix}/")) { + // MCP-required OpenID Connect path-insertion compatibility form + suffix + } else if let Some(prefix) = path.strip_suffix(oidc_prefix) { + // OpenID Connect path-appended form + prefix.trim_start_matches('/') + } else { + return None; + }; + + let mut issuer = discovery_url.clone(); + issuer.set_query(None); + issuer.set_fragment(None); + if issuer_path.is_empty() { + issuer.set_path(""); + } else { + issuer.set_path(&format!("/{issuer_path}")); + } + Some(issuer.to_string()) + } + + fn issuer_identifiers_match(received_issuer: &str, expected_issuer: &str) -> bool { + if received_issuer == expected_issuer { + return true; + } + + let trim_root_slash = |issuer: &str| -> String { + issuer + .strip_suffix('/') + .filter(|without_slash| { + Url::parse(without_slash) + .map(|url| url.path().is_empty() || url.path() == "/") + .unwrap_or(false) + }) + .unwrap_or(issuer) + .to_string() + }; + + trim_root_slash(received_issuer) == trim_root_slash(expected_issuer) + } + + fn validate_authorization_metadata_issuer( + &self, + discovery_url: &Url, + expected_issuer: Option<&str>, + metadata: &AuthorizationMetadata, + ) -> Result<(), AuthError> { + let expected_issuer = expected_issuer + .map(str::to_owned) + .or_else(|| Self::expected_issuer_for_authorization_metadata_url(discovery_url)); + let Some(expected_issuer) = expected_issuer else { + return Ok(()); + }; + let Some(received_issuer) = metadata.issuer.as_deref() else { + if self.allow_missing_issuer { + return Ok(()); + } + return Err(AuthError::AuthorizationServerMissingIssuer { expected_issuer }); + }; + if !Self::issuer_identifiers_match(received_issuer, &expected_issuer) { + return Err(AuthError::AuthorizationServerMismatch { + expected_issuer, + received_issuer: received_issuer.to_string(), + }); + } + Ok(()) + } + async fn discover_oauth_server_via_resource_metadata( &self, ) -> Result, AuthError> { - let Some(resource_metadata_url) = self.discover_resource_metadata_url().await? else { + let Some((resource_metadata_url, resource_metadata)) = + self.discover_resource_metadata().await? + else { return Ok(None); }; + self.authorization_metadata_from_resource_metadata( + &resource_metadata_url, + resource_metadata, + ) + .await + } + /// Read protected resource metadata from the url a `WWW-Authenticate` + /// challenge advertised. + async fn discover_oauth_server_from_resource_metadata_url( + &self, + resource_metadata_url: &Url, + ) -> Result, AuthError> { let Some(resource_metadata) = self - .fetch_resource_metadata_from_url(&resource_metadata_url) + .fetch_resource_metadata_from_url( + resource_metadata_url, + ResourceMetadataUrlOrigin::Advertised, + ) .await? else { return Ok(None); }; - self.validate_resource_metadata_resource(&resource_metadata)?; + self.authorization_metadata_from_resource_metadata(resource_metadata_url, resource_metadata) + .await + } + + /// Walk the authorization servers a protected resource metadata document + /// names, keeping the first one that answers with usable metadata. A document + /// that names servers and none of them answers is an error, not a reason to + /// go looking elsewhere; a document that names none returns `Ok(None)`. + /// + /// The document arrives here through `read_resource_metadata`, which is where + /// it is decided to be this resource's metadata at all. + async fn authorization_metadata_from_resource_metadata( + &self, + resource_metadata_url: &Url, + resource_metadata: ResourceServerMetadata, + ) -> Result, AuthError> { + self.discovered_resource + .write() + .await + .replace(resource_metadata.resource.clone().unwrap_or_default()); // store scopes_supported from protected resource metadata for select_scopes() - if let Some(scopes) = resource_metadata.scopes_supported { - if !scopes.is_empty() { - *self.resource_scopes.write().await = scopes; - } + if let Some(scopes) = resource_metadata.scopes_supported + && !scopes.is_empty() + { + *self.resource_scopes.write().await = scopes; } - let mut candidates = Vec::new(); + // A server naming the same authorization server in both the singular draft + // field and the list would otherwise have each of that server's well-known + // forms requested twice. + let mut candidates: Vec = Vec::new(); + let mut push_candidate = |candidate: String| { + let candidate = candidate.trim(); + if !candidate.is_empty() && !candidates.iter().any(|kept| kept == candidate) { + candidates.push(candidate.to_string()); + } + }; if let Some(single) = resource_metadata.authorization_server { - candidates.push(single); + push_candidate(single); } - if let Some(list) = resource_metadata.authorization_servers { - candidates.extend(list); + for candidate in resource_metadata.authorization_servers.unwrap_or_default() { + push_candidate(candidate); } - for candidate in candidates { - let candidate = candidate.trim(); - if candidate.is_empty() { - continue; - } + if candidates.is_empty() { + debug!( + "protected resource metadata at {resource_metadata_url} names no authorization server" + ); + return Ok(None); + } + for candidate in &candidates { let candidate_url = match Url::parse(candidate) { Ok(url) => url, Err(_) => match resource_metadata_url.join(candidate) { @@ -1882,24 +2631,37 @@ impl AuthorizationManager { }, }; - if !Self::is_allowed_authorization_server_metadata_url(&candidate_url) { + if !Self::is_allowed_authorization_server_metadata_url(&self.base_url, &candidate_url) { warn!("rejecting authorization server metadata URL `{candidate_url}`"); continue; } if candidate_url.path().contains("/.well-known/") { - if let Some(metadata) = self.fetch_authorization_metadata(&candidate_url).await? { + if let Some(metadata) = self + .fetch_authorization_metadata(&candidate_url, None) + .await? + { return Ok(Some(metadata)); } continue; } - if let Some(metadata) = self.try_discover_oauth_server(&candidate_url).await? { + // Discovery URL construction removes a non-root trailing slash. Keep the issuer + // advertised by protected resource metadata for the exact RFC 8414 comparison. + if let Some(metadata) = self + .try_discover_oauth_server(&candidate_url, Some(candidate_url.as_str())) + .await? + { return Ok(Some(metadata)); } } - Ok(None) + // Falling back to the base url would send the user somewhere the resource + // never named. + Err(AuthError::MetadataError(format!( + "protected resource metadata at {resource_metadata_url} names authorization servers {}, but none published usable metadata", + candidates.join(", ") + ))) } fn validate_resource_metadata_resource( @@ -1912,9 +2674,21 @@ impl AuthorizationManager { )); }; - if !Self::resource_identifiers_match(self.base_url.as_str(), resource) { + let Ok(resource_url) = Url::parse(resource) else { + return Err(AuthError::MetadataError( + "Protected resource metadata resource field is not a valid URL".to_string(), + )); + }; + + if resource_url.fragment().is_some() { + return Err(AuthError::MetadataError( + "Protected resource metadata resource does not permit fragment in URL as specified by RFC 8707".to_string() + )); + } + + if !Self::is_resource_identifier_valid(&self.base_url, &resource_url) { return Err(AuthError::MetadataError(format!( - "Protected resource metadata resource mismatch: expected '{}', got '{}'", + "Protected resource metadata resource mismatch: reference '{}', permitted '{}'", self.base_url, resource ))); } @@ -1922,102 +2696,175 @@ impl AuthorizationManager { Ok(()) } - fn resource_identifiers_match(expected: &str, actual: &str) -> bool { - expected == actual - || (Self::is_root_resource_identifier(expected) - && actual == expected.trim_end_matches('/')) - || (Self::is_root_resource_identifier(actual) - && expected == actual.trim_end_matches('/')) - } - - fn is_root_resource_identifier(value: &str) -> bool { - Url::parse(value) - .is_ok_and(|url| url.path() == "/" && url.query().is_none() && url.fragment().is_none()) - } - - async fn discover_resource_metadata_url(&self) -> Result, AuthError> { - if let Ok(Some(resource_metadata_url)) = - self.fetch_resource_metadata_url(&self.base_url).await - { - return Ok(Some(resource_metadata_url)); + fn is_resource_identifier_valid(expected: &Url, actual: &Url) -> bool { + if expected == actual { + return true; } - // If the primary URL doesn't use WWW-Authenticate, try oauth-protected-resource discovery. + if expected.scheme() != actual.scheme() + || expected.host_str() != actual.host_str() + || expected.port_or_known_default() != actual.port_or_known_default() + { + return false; + } + + let expected_path = expected.path(); + let actual_path = actual.path(); + + // Query parameters may carry transport hints rather than resource identity. + if expected_path == actual_path { + return true; + } + + expected_path.starts_with(actual_path) + && (actual_path.ends_with('/') + || expected_path.as_bytes().get(actual_path.len()) == Some(&b'/')) + } + + /// Look for the protected resource metadata document, reading it where it is + /// found so that a candidate answering with something else only costs that + /// candidate. + async fn discover_resource_metadata( + &self, + ) -> Result, AuthError> { + // A url the resource points at can also be one of the candidates below. + let mut requested = HashSet::new(); + + if let Some(advertised_url) = self.probe_resource_endpoint_for_challenge().await? { + requested.insert(advertised_url.clone()); + if let Some(metadata) = self + .fetch_resource_metadata_from_url( + &advertised_url, + ResourceMetadataUrlOrigin::Advertised, + ) + .await? + { + return Ok(Some((advertised_url, metadata))); + } + // Nothing was published there. The candidates below are reached from the + // base url rather than from that pointer, so they are still worth trying. + } + + // The other place the document can be is the well-known location. // https://www.rfc-editor.org/rfc/rfc9728.html#name-obtaining-protected-resourc for candidate_path in Self::well_known_paths(self.base_url.path(), "oauth-protected-resource") { - let mut discovery_url = self.base_url.clone(); - discovery_url.set_query(None); - discovery_url.set_fragment(None); - discovery_url.set_path(&candidate_path); - if let Ok(Some(resource_metadata_url)) = - self.fetch_resource_metadata_url(&discovery_url).await - { - return Ok(Some(resource_metadata_url)); + let mut candidate_url = self.base_url.clone(); + candidate_url.set_query(None); + candidate_url.set_fragment(None); + candidate_url.set_path(&candidate_path); + + if !requested.insert(candidate_url.clone()) { + continue; + } + + let response = self + .discovery_get(&candidate_url) + .await + .map_err(|error| Self::discovery_failed(&candidate_url, error))?; + + match response.status() { + // The candidate url is the document itself, so read the body here + // instead of requesting the same url again. + StatusCode::OK => { + if let Some(metadata) = self.read_resource_metadata( + &candidate_url, + response.body(), + ResourceMetadataUrlOrigin::WellKnownGuess, + )? { + return Ok(Some((candidate_url, metadata))); + } + } + StatusCode::UNAUTHORIZED => { + let Some(pointer_url) = self + .extract_resource_metadata_url_from_www_authenticate(&response) + .await + else { + continue; + }; + if !requested.insert(pointer_url.clone()) { + continue; + } + // Found on a url this run guessed at, so a wrong answer only rules + // out the pointer. + if let Some(metadata) = self + .fetch_resource_metadata_from_url( + &pointer_url, + ResourceMetadataUrlOrigin::WellKnownGuess, + ) + .await? + { + return Ok(Some((pointer_url, metadata))); + } + } + status => debug!("resource metadata probe returned unexpected status: {status}"), } } Ok(None) } - /// Extract the resource metadata url from the WWW-Authenticate header value. + /// Probe the resource itself, looking only for a `WWW-Authenticate` challenge + /// that carries a `resource_metadata` pointer. + /// + /// A 200 here says nothing about metadata. RFC 9728 publishes the document at + /// the well-known URI and advertises it through the challenge parameter, so the + /// resource answering its own GET is not the document and must not end + /// discovery before the well-known candidates are tried. /// https://www.rfc-editor.org/rfc/rfc9728.html#name-use-of-www-authenticate-for - async fn fetch_resource_metadata_url(&self, url: &Url) -> Result, AuthError> { - let response = match self.discovery_get(url).await { - Ok(r) => r, - Err(e) => { - debug!("resource metadata probe failed: {}", e); - return Ok(None); - } - }; + async fn probe_resource_endpoint_for_challenge(&self) -> Result, AuthError> { + let response = self + .discovery_get(&self.base_url) + .await + .map_err(|error| Self::discovery_failed(&self.base_url, error))?; - if response.status() == StatusCode::OK { - return Ok(Some(url.clone())); - } else if response.status() != StatusCode::UNAUTHORIZED { - debug!( - "resource metadata probe returned unexpected status: {}", - response.status() - ); - return Ok(None); + if response.status() == StatusCode::UNAUTHORIZED { + return Ok(self + .extract_resource_metadata_url_from_www_authenticate(&response) + .await); } + debug!( + "resource endpoint probe returned {}, no WWW-Authenticate pointer to follow", + response.status() + ); + Ok(None) + } + + async fn extract_resource_metadata_url_from_www_authenticate( + &self, + response: &HttpResponse, + ) -> Option { let mut parsed_url = None; for value in response.headers().get_all(WWW_AUTHENTICATE).iter() { let Ok(value_str) = value.to_str() else { continue; }; let params = Self::extract_www_authenticate_params(value_str, &self.base_url); - if let Some(url) = params.resource_metadata_url { - if let Some(scope) = ¶ms.scope { - debug!("WWW-Authenticate header contains scope: {}", scope); - let scopes: Vec = - scope.split_whitespace().map(|s| s.to_string()).collect(); - *self.www_auth_scopes.write().await = scopes; - } - parsed_url = Some(url); + if params.resource_metadata_url.is_some() { + self.record_challenge_scope(¶ms).await; + parsed_url = params.resource_metadata_url; break; } } - Ok(parsed_url) + parsed_url } async fn fetch_resource_metadata_from_url( &self, resource_metadata_url: &Url, + origin: ResourceMetadataUrlOrigin, ) -> Result, AuthError> { debug!( "resource metadata discovery url: {:?}", resource_metadata_url ); - let response = match self.discovery_get(resource_metadata_url).await { - Ok(r) => r, - Err(e) => { - debug!("resource metadata request failed: {}", e); - return Ok(None); - } - }; + let response = self + .discovery_get(resource_metadata_url) + .await + .map_err(|error| Self::discovery_failed(resource_metadata_url, error))?; if response.status() != StatusCode::OK { debug!( @@ -2027,16 +2874,96 @@ impl AuthorizationManager { return Ok(None); } - let metadata = match serde_json::from_slice::(response.body()) { + self.read_resource_metadata(resource_metadata_url, response.body(), origin) + } + + /// Read a response body as this resource's protected resource metadata. + /// + /// A body that is not that document rules out the url it came from, and where + /// that url came from decides whether ruling it out leaves anything to try. + fn read_resource_metadata( + &self, + resource_metadata_url: &Url, + body: &[u8], + origin: ResourceMetadataUrlOrigin, + ) -> Result, AuthError> { + let metadata = match serde_json::from_slice::(body) { Ok(metadata) => metadata, Err(e) => { debug!("failed to parse resource metadata as JSON: {}", e); return Ok(None); } }; + + // Every field of `ResourceServerMetadata` is optional, so an unrelated JSON + // object deserializes into an all-`None` value and then fails validation + // fatally. RFC 9728 requires `resource`, and MCP requires an authorization + // server reference, so a document carrying neither is not a protected + // resource metadata document. + if metadata.resource.is_none() + && metadata.authorization_server.is_none() + && metadata.authorization_servers.is_none() + { + return match origin { + // The server named this url, so there is nothing better to move on + // to: the alternatives all drop the resource binding the document + // was supposed to carry. Report it instead. + ResourceMetadataUrlOrigin::Advertised => Err(AuthError::MetadataError(format!( + "the server advertised {resource_metadata_url} as protected resource metadata, but the document carries neither `resource` nor an authorization server reference" + ))), + // Nothing advertised this url, so its answer only rules out this + // candidate. + ResourceMetadataUrlOrigin::WellKnownGuess => { + debug!( + "response at {resource_metadata_url} is not a protected resource metadata document" + ); + Ok(None) + } + }; + } + + // Carrying those fields only makes the body a metadata document; validation + // is what makes it this resource's. Both answers rule out the url the same + // way, so both are read the same way. + if let Err(error) = self.validate_resource_metadata_resource(&metadata) { + return match origin { + ResourceMetadataUrlOrigin::Advertised => Err(error), + ResourceMetadataUrlOrigin::WellKnownGuess => { + debug!( + "document at {resource_metadata_url} is not this resource's metadata: {error}" + ); + Ok(None) + } + }; + } + Ok(Some(metadata)) } + fn discovery_failed(url: &Url, error: OAuthHttpClientError) -> AuthError { + AuthError::MetadataError(format!( + "OAuth metadata discovery failed for {url}\n Caused by: {}", + crate::error::ErrorChain(error.as_ref()) + )) + } + + async fn discovery_request( + &self, + request: OAuthHttpRequest, + ) -> Result { + let response = self.http_client.execute(request).await?; + let status = response.status(); + if status.is_server_error() + || matches!( + status, + StatusCode::REQUEST_TIMEOUT | StatusCode::TOO_EARLY | StatusCode::TOO_MANY_REQUESTS + ) + { + return Err(Box::new(OAuthHttpError::UnexpectedStatus(status))); + } + Ok(response) + } + async fn discovery_get(&self, url: &Url) -> Result { let mut current_url = url.clone(); for _ in 0..MAX_OAUTH_DISCOVERY_REDIRECTS { @@ -2045,10 +2972,9 @@ impl AuthorizationManager { .uri(current_url.as_str()) .header(HEADER_MCP_PROTOCOL_VERSION, "2024-11-05") .body(Vec::new()) - .map_err(|error| OAuthHttpClientError::new(error.to_string()))?; + .map_err(|error| Box::new(error) as OAuthHttpClientError)?; let response = self - .http_client - .execute(OAuthHttpRequest::new( + .discovery_request(OAuthHttpRequest::new( request, OAuthHttpRedirectPolicy::Stop, )) @@ -2063,23 +2989,21 @@ impl AuthorizationManager { }; let location = location .to_str() - .map_err(|error| OAuthHttpClientError::new(error.to_string()))?; + .map_err(|error| Box::new(error) as OAuthHttpClientError)?; let next_url = current_url .join(location) - .map_err(|error| OAuthHttpClientError::new(error.to_string()))?; + .map_err(|error| Box::new(error) as OAuthHttpClientError)?; if Self::is_http_url(&next_url) && Self::is_same_origin(¤t_url, &next_url) { current_url = next_url; continue; } - return Err(OAuthHttpClientError::new(format!( - "OAuth discovery redirect to non-same-origin URL rejected: {next_url}" - ))); + return Err(Box::new(OAuthHttpError::CrossOriginRedirect(next_url))); } - Err(OAuthHttpClientError::new(format!( - "OAuth discovery exceeded {MAX_OAUTH_DISCOVERY_REDIRECTS} redirects" + Err(Box::new(OAuthHttpError::TooManyRedirects( + MAX_OAUTH_DISCOVERY_REDIRECTS, ))) } @@ -2227,20 +3151,18 @@ impl AuthorizationManager { if let ClientCredentialsConfig::PrivateKeyJwt { signing_algorithm, .. } = config - { - if let Some(algs) = metadata + && let Some(algs) = metadata .additional_fields .get("token_endpoint_auth_signing_alg_values_supported") .and_then(|v| v.as_array()) - { - let alg_str = signing_algorithm.as_str(); - if !algs.iter().any(|a| a.as_str() == Some(alg_str)) { - let supported: Vec<&str> = algs.iter().filter_map(|a| a.as_str()).collect(); - return Err(AuthError::ClientCredentialsError(format!( - "Authorization server does not support signing algorithm '{}'. Supported: {:?}", - alg_str, supported - ))); - } + { + let alg_str = signing_algorithm.as_str(); + if !algs.iter().any(|a| a.as_str() == Some(alg_str)) { + let supported: Vec<&str> = algs.iter().filter_map(|a| a.as_str()).collect(); + return Err(AuthError::ClientCredentialsError(format!( + "Authorization server does not support signing algorithm '{}'. Supported: {:?}", + alg_str, supported + ))); } } @@ -2397,6 +3319,7 @@ impl AuthorizationManager { token_response: Some(token_result.clone()), granted_scopes, token_received_at: Some(Self::now_epoch_secs()), + issuer: self.metadata_issuer(), }; self.credential_store.save(stored).await?; @@ -2431,22 +3354,20 @@ impl AuthorizationManager { .as_ref() .ok_or(AuthError::NoAuthorizationSupport)?; - // Validate that the token endpoint uses HTTPS before transmitting sensitive credentials. let token_endpoint_url = url::Url::parse(&metadata.token_endpoint).map_err(|e| { AuthError::ClientCredentialsError(format!( "Invalid token endpoint URL in authorization metadata: {e}" )) })?; - if token_endpoint_url.scheme() != "https" { + if !is_allowed_client_credentials_endpoint(&self.base_url, &token_endpoint_url) { return Err(AuthError::ClientCredentialsError( "Insecure token endpoint URL: HTTPS is required for client credentials flow" .to_string(), )); } - let audience = token_endpoint_audience - .as_deref() - .unwrap_or(&metadata.token_endpoint); + let audience = + client_authentication_audience(metadata, token_endpoint_audience.as_deref())?; let assertion = Self::build_jwt_assertion(client_id, audience, signing_key, *signing_algorithm)?; @@ -2519,6 +3440,7 @@ impl AuthorizationManager { token_response: Some(token_result.clone()), granted_scopes, token_received_at: Some(Self::now_epoch_secs()), + issuer: self.metadata_issuer(), }; self.credential_store.save(stored).await?; @@ -2551,7 +3473,7 @@ impl AuthorizationManager { "jti": jti, }); - let header = jsonwebtoken::Header::new(algorithm.to_jsonwebtoken_algorithm()); + let header = client_authentication_header(algorithm); let encoding_key = jsonwebtoken::EncodingKey::from_rsa_pem(signing_key).or_else(|_| { jsonwebtoken::EncodingKey::from_ec_pem(signing_key).map_err(|e| { AuthError::JwtSigningError(format!("Failed to parse signing key: {}", e)) @@ -2613,16 +3535,56 @@ pub struct AuthorizationSession { } impl AuthorizationSession { - /// create new authorization session + /// Create a new authorization session, selecting a client registration + /// mechanism per the [MCP authorization specification](https://modelcontextprotocol.io/specification/2026-07-28/basic/authorization/client-registration) + /// priority order: + /// + /// 1. Pre-registered client information + /// ([`AuthorizationRequest::with_preregistered_client`]), when available + /// 2. Client ID Metadata Documents + /// ([`AuthorizationRequest::with_client_metadata_url`]), when the + /// authorization server advertises `client_id_metadata_document_supported` + /// 3. Dynamic Client Registration, when the authorization server + /// advertises a `registration_endpoint` + /// + /// The manager must already have discovered authorization server metadata. + /// If `request.scopes` is empty, scopes are selected using the SDK's + /// normal scope-selection policy. + /// + /// On failure, the manager is returned alongside the error so callers can + /// retry without losing the original configuration and stores. pub async fn new( mut auth_manager: AuthorizationManager, - scopes: &[&str], - redirect_uri: &str, - client_name: Option<&str>, - client_metadata_url: Option<&str>, - ) -> Result { - let metadata = auth_manager.metadata.as_ref(); - let supports_url_based_client_id = metadata + mut request: AuthorizationRequest, + ) -> Result { + if request.client_secret.is_some() && request.client_id.is_none() { + return Err(( + auth_manager, + AuthError::RegistrationFailed( + "client_secret was provided without a pre-registered client_id; \ + pair with_client_secret with with_preregistered_client" + .to_string(), + ), + )); + } + + if request.scopes.is_empty() { + request.scopes = auth_manager.select_scopes(None, &[]); + } else { + auth_manager.add_offline_access_if_supported(&mut request.scopes); + } + + if request.application_type.is_some() { + auth_manager.application_type = request.application_type.clone(); + } + + let redirect_uri = request.redirect_uri.clone(); + let scopes = request.scopes.clone(); + let scope_refs: Vec<&str> = scopes.iter().map(|s| s.as_str()).collect(); + + let supports_url_based_client_id = auth_manager + .metadata + .as_ref() .and_then(|m| { m.additional_fields .get("client_id_metadata_document_supported") @@ -2630,56 +3592,78 @@ impl AuthorizationSession { .and_then(|v| v.as_bool()) .unwrap_or(false); - let config = if supports_url_based_client_id { - if let Some(client_metadata_url) = client_metadata_url { - if !is_https_url(client_metadata_url) { - return Err(AuthError::RegistrationFailed(format!( + // 1. pre-registered client information takes priority over everything else + let config = if let Some(client_id) = &request.client_id { + OAuthClientConfig { + client_id: client_id.clone(), + client_secret: request.client_secret.clone(), + scopes: scopes.clone(), + redirect_uri: redirect_uri.clone(), + application_type: request.application_type.clone(), + } + // 2. CIMD (SEP-991), when the server advertises support and the client hosts a metadata document + } else if let Some(client_metadata_url) = request.client_metadata_url.as_deref() + && supports_url_based_client_id + { + if !is_https_url(client_metadata_url) { + return Err(( + auth_manager, + AuthError::RegistrationFailed(format!( "client_metadata_url must be a valid HTTPS URL with a non-root pathname, got: {}", client_metadata_url - ))); - } - // SEP-991: URL-based Client IDs - use URL as client_id directly. - // SEP-837: match the hosted client-metadata.json application_type ("native") - OAuthClientConfig { - client_id: client_metadata_url.to_string(), - client_secret: None, - scopes: scopes.iter().map(|s| s.to_string()).collect(), - redirect_uri: redirect_uri.to_string(), - application_type: Some(DEFAULT_APPLICATION_TYPE.to_string()), - } - } else { - // Fallback to dynamic registration - auth_manager - .register_client(client_name.unwrap_or("MCP Client"), redirect_uri, scopes) - .await - .map_err(|e| { - AuthError::RegistrationFailed(format!("Dynamic registration failed: {}", e)) - })? + )), + )); } + // SEP-991: URL-based Client IDs - use URL as client_id directly. + // SEP-837: match the hosted client-metadata.json application_type ("native") + OAuthClientConfig { + client_id: client_metadata_url.to_string(), + client_secret: None, + scopes: scopes.clone(), + redirect_uri: redirect_uri.clone(), + application_type: Some( + request + .application_type + .clone() + .unwrap_or_else(|| DEFAULT_APPLICATION_TYPE.to_string()), + ), + } + // 3. fall back to dynamic client registration } else { - // Fallback to dynamic registration match auth_manager - .register_client(client_name.unwrap_or("MCP Client"), redirect_uri, scopes) + .register_client( + request.client_name.as_deref().unwrap_or("MCP Client"), + &redirect_uri, + &scope_refs, + ) .await { Ok(config) => config, Err(e) => { - return Err(AuthError::RegistrationFailed(format!( - "Dynamic registration failed: {}", - e - ))); + return Err(( + auth_manager, + AuthError::RegistrationFailed(format!( + "Dynamic registration failed: {}", + e + )), + )); } } }; // reset client config - auth_manager.configure_client(config)?; - let auth_url = auth_manager.get_authorization_url(scopes).await?; + if let Err(e) = auth_manager.configure_client(config) { + return Err((auth_manager, e)); + } + let auth_url = match auth_manager.get_authorization_url(&scope_refs).await { + Ok(url) => url, + Err(e) => return Err((auth_manager, e)), + }; Ok(Self { auth_manager, auth_url, - redirect_uri: redirect_uri.to_string(), + redirect_uri, }) } @@ -2811,7 +3795,7 @@ impl OAuthState { ) } - async fn placeholder(&self) -> Result { + async fn placeholder_state(&self) -> Result { let (http_client, refresh_redirect_policy) = self.oauth_http_client_config(); Ok(OAuthState::Unauthorized( AuthorizationManager::new_inner( @@ -2881,17 +3865,18 @@ impl OAuthState { *manager.current_scopes.write().await = granted_scopes.clone(); + let resolution = manager.resolve_metadata().await?; + manager.metadata = Some(resolution.metadata); + let stored = StoredCredentials { client_id: client_id.to_string(), token_response: Some(credentials), granted_scopes, token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: manager.metadata_issuer(), }; manager.credential_store.save(stored).await?; - let metadata = manager.discover_metadata().await?; - manager.metadata = Some(metadata); - manager.configure_client_id(client_id)?; *self = OAuthState::Authorized(manager); @@ -2903,59 +3888,65 @@ impl OAuthState { } } - /// start authorization + /// Start authorization. + /// + /// Selects a client registration mechanism from the identity material in + /// `request`, following the [MCP authorization specification](https://modelcontextprotocol.io/specification/2026-07-28/basic/authorization/client-registration) + /// priority order: + /// + /// 1. Pre-registered client information + /// ([`AuthorizationRequest::with_preregistered_client`]), when available + /// 2. Client ID Metadata Documents + /// ([`AuthorizationRequest::with_client_metadata_url`]), when the + /// authorization server advertises `client_id_metadata_document_supported` + /// 3. Dynamic Client Registration, when the authorization server + /// advertises a `registration_endpoint` + /// + /// If `request.scopes` is empty, scopes are selected using the SDK's + /// normal scope-selection policy. + /// + /// On failure, the state returns to `Unauthorized` so callers can retry + /// without losing the original configuration and stores. pub async fn start_authorization( &mut self, - scopes: &[&str], - redirect_uri: &str, - client_name: Option<&str>, + request: AuthorizationRequest, ) -> Result<(), AuthError> { - self.start_authorization_with_metadata_url(scopes, redirect_uri, client_name, None) - .await - } - - /// start authorization with optional client metadata URL (SEP-991) - pub async fn start_authorization_with_metadata_url( - &mut self, - scopes: &[&str], - redirect_uri: &str, - client_name: Option<&str>, - client_metadata_url: Option<&str>, - ) -> Result<(), AuthError> { - let placeholder = self.placeholder().await?; - if let OAuthState::Unauthorized(mut manager) = std::mem::replace(self, placeholder) { - debug!("start discovery"); - let metadata = manager.discover_metadata().await?; - manager.metadata = Some(metadata); - let selected_scopes: Vec = if scopes.is_empty() { - manager.select_scopes(None, &[]) - } else { - let mut s: Vec = scopes.iter().map(|s| s.to_string()).collect(); - manager.add_offline_access_if_supported(&mut s); - s - }; - let scope_refs: Vec<&str> = selected_scopes.iter().map(|s| s.as_str()).collect(); - debug!("start session"); - let session = AuthorizationSession::new( - manager, - &scope_refs, - redirect_uri, - client_name, - client_metadata_url, - ) - .await?; - *self = OAuthState::Session(session); - Ok(()) - } else { - Err(AuthError::InternalError( + let placeholder = self.placeholder_state().await?; + let old = std::mem::replace(self, placeholder); + let OAuthState::Unauthorized(mut manager) = old else { + *self = old; + return Err(AuthError::InternalError( "Already in session state".to_string(), - )) + )); + }; + debug!("start discovery"); + let resolution = manager + .resolve_metadata_from_challenge(request.challenge.as_deref()) + .await; + let metadata = match resolution { + Ok(resolution) => resolution.metadata, + Err(e) => { + *self = OAuthState::Unauthorized(manager); + return Err(e); + } + }; + manager.metadata = Some(metadata); + debug!("start session"); + match AuthorizationSession::new(manager, request).await { + Ok(session) => { + *self = OAuthState::Session(session); + Ok(()) + } + Err((manager, e)) => { + *self = OAuthState::Unauthorized(manager); + Err(e) + } } } /// complete authorization pub async fn complete_authorization(&mut self) -> Result<(), AuthError> { - let placeholder = self.placeholder().await?; + let placeholder = self.placeholder_state().await?; if let OAuthState::Session(session) = std::mem::replace(self, placeholder) { *self = OAuthState::Authorized(session.auth_manager); Ok(()) @@ -2963,9 +3954,9 @@ impl OAuthState { Err(AuthError::InternalError("Not in session state".to_string())) } } - /// covert to authorized http client + /// convert to authorized http client pub async fn to_authorized_http_client(&mut self) -> Result<(), AuthError> { - let placeholder = self.placeholder().await?; + let placeholder = self.placeholder_state().await?; if let OAuthState::Authorized(manager) = std::mem::replace(self, placeholder) { *self = OAuthState::AuthorizedHttpClient(AuthorizedHttpClient::new( Arc::new(manager), @@ -2985,7 +3976,7 @@ impl OAuthState { required_scope: &str, redirect_uri: &str, ) -> Result { - let placeholder = self.placeholder().await?; + let placeholder = self.placeholder_state().await?; let old = std::mem::replace(self, placeholder); let OAuthState::Authorized(manager) = old else { *self = old; @@ -3117,7 +4108,7 @@ impl OAuthState { &mut self, config: ClientCredentialsConfig, ) -> Result<(), AuthError> { - let placeholder = self.placeholder().await?; + let placeholder = self.placeholder_state().await?; let OAuthState::Unauthorized(mut manager) = std::mem::replace(self, placeholder) else { return Err(AuthError::InternalError( "Client credentials flow requires Unauthorized state".to_string(), @@ -3125,8 +4116,8 @@ impl OAuthState { }; // Discover metadata - let metadata = manager.discover_metadata().await?; - manager.metadata = Some(metadata); + let resolution = manager.resolve_metadata().await?; + manager.metadata = Some(resolution.metadata); // Validate server supports the requested auth method manager.validate_client_credentials_metadata(&config)?; @@ -3149,15 +4140,19 @@ mod tests { sync::{Arc, Mutex as StdMutex}, }; - use oauth2::{AuthType, CsrfToken, HttpResponse, PkceCodeVerifier}; + use oauth2::{AuthType, CsrfToken, HttpResponse, PkceCodeVerifier, TokenResponse}; + use reqwest::StatusCode; use rstest::rstest; + use tokio::sync::{Mutex, OwnedMutexGuard, Semaphore}; use url::Url; use super::{ AuthError, AuthorizationCallback, AuthorizationManager, AuthorizationMetadata, - InMemoryStateStore, OAuthClientConfig, OAuthHttpClient, OAuthHttpClientError, - OAuthHttpClientFuture, OAuthHttpRedirectPolicy, OAuthHttpRequest, ScopeUpgradeConfig, - StateStore, StoredAuthorizationState, is_https_url, + AuthorizationMetadataSource, AuthorizationRequest, AuthorizationSession, + CredentialRefreshGuard, CredentialStore, InMemoryCredentialStore, InMemoryStateStore, + OAuthClientConfig, OAuthHttpClient, OAuthHttpClientError, OAuthHttpClientFuture, + OAuthHttpRedirectPolicy, OAuthHttpRequest, ScopeUpgradeConfig, StateStore, + StoredAuthorizationState, is_https_url, }; use crate::transport::auth::VendorExtraTokenFields; @@ -3197,9 +4192,7 @@ mod tests { body: request.request.body().clone(), }); let response = self.responses.lock().unwrap().pop_front(); - Box::pin(async move { - response.ok_or_else(|| OAuthHttpClientError::new("missing fake response")) - }) + Box::pin(async move { response.ok_or_else(|| "missing fake response".into()) }) } } @@ -3218,6 +4211,226 @@ mod tests { .unwrap() } + fn empty_response(status: u16) -> HttpResponse { + oauth2::http::Response::builder() + .status(status) + .body(Vec::new()) + .unwrap() + } + + #[test] + fn oauth_http_client_error_preserves_source_chain() { + #[derive(Debug, thiserror::Error)] + #[error("request failed")] + struct RequestError(#[source] std::io::Error); + + let error: OAuthHttpClientError = RequestError(std::io::Error::other( + "certificate signed by unknown authority", + )) + .into(); + assert!(error.downcast_ref::().is_some()); + + let url = Url::parse("https://mcp.example.com/mcp").unwrap(); + let error = AuthorizationManager::discovery_failed(&url, error); + assert_eq!( + error.to_string(), + "Metadata error: OAuth metadata discovery failed for https://mcp.example.com/mcp\n Caused by: request failed\n Caused by: certificate signed by unknown authority" + ); + } + + #[tokio::test] + async fn default_oauth_http_client_honors_redirect_policy() { + use axum::{Router, routing::post}; + + let received = Arc::new(StdMutex::new(Vec::new())); + let capture = Arc::clone(&received); + let app = Router::new() + .route( + "/redirect", + post(|| async { (StatusCode::TEMPORARY_REDIRECT, [("location", "/token")]) }), + ) + .route( + "/token", + post(move |body: String| { + let capture = Arc::clone(&capture); + async move { + capture.lock().unwrap().push(body); + StatusCode::OK + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}/redirect", listener.local_addr().unwrap()); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let client = super::default_oauth_http_client().unwrap(); + + for (policy, expected) in [ + ( + OAuthHttpRedirectPolicy::Stop, + StatusCode::TEMPORARY_REDIRECT, + ), + (OAuthHttpRedirectPolicy::Follow, StatusCode::OK), + ] { + let request = oauth2::http::Request::builder() + .method("POST") + .uri(&endpoint) + .body(b"credential-sentinel".to_vec()) + .unwrap(); + let response = client + .execute(OAuthHttpRequest::new(request, policy)) + .await + .unwrap(); + assert_eq!(response.status(), expected); + let expected_bodies = if policy == OAuthHttpRedirectPolicy::Stop { + vec![] + } else { + vec!["credential-sentinel".to_owned()] + }; + assert_eq!(*received.lock().unwrap(), expected_bodies); + } + } + + #[tokio::test] + async fn default_http_client_preserves_connection_failure_cause() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}/mcp", listener.local_addr().unwrap()); + drop(listener); + + let manager = AuthorizationManager::new(&url).await.unwrap(); + let error = manager.resolve_metadata().await.unwrap_err(); + + assert!( + matches!( + error, + AuthError::MetadataError(ref reason) + if reason.contains(&url) + && reason.contains("\n Caused by: error sending request for url") + && reason.matches("error sending request for url").count() == 1 + && reason.to_ascii_lowercase().contains("connection refused") + ), + "unexpected discovery error: {error}" + ); + } + + #[tokio::test] + async fn authorization_metadata_propagates_transport_failure() { + let responses = preregistered_discovery_responses() + .into_iter() + .take(2) + .collect(); + let client = RecordingOAuthHttpClient::with_responses(responses); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let error = manager.resolve_metadata().await.unwrap_err(); + + assert!( + matches!( + error, + AuthError::MetadataError(ref reason) + if reason.contains("https://auth.example.com/.well-known/oauth-authorization-server") + && reason.contains("missing fake response") + ), + "unexpected discovery error: {error}" + ); + assert_eq!(client.requests().len(), 3); + } + + #[tokio::test] + async fn discovery_propagates_server_errors() { + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(RecordingOAuthHttpClient::with_responses(vec![ + empty_response(503), + ])), + ) + .await + .unwrap(); + + let error = manager.resolve_metadata().await.unwrap_err(); + + assert!( + matches!( + error, + AuthError::MetadataError(ref reason) + if reason.contains("https://mcp.example.com/mcp") && reason.contains("503") + ), + "unexpected discovery error: {error}" + ); + } + + #[rstest] + #[case::resource_request_timeout(StatusCode::REQUEST_TIMEOUT, 0, "https://mcp.example.com/mcp")] + #[case::resource_too_early(StatusCode::TOO_EARLY, 0, "https://mcp.example.com/mcp")] + #[case::resource_too_many_requests( + StatusCode::TOO_MANY_REQUESTS, + 0, + "https://mcp.example.com/mcp" + )] + #[case::protected_metadata_request_timeout( + StatusCode::REQUEST_TIMEOUT, + 1, + "https://mcp.example.com/.well-known/oauth-protected-resource" + )] + #[case::protected_metadata_too_early( + StatusCode::TOO_EARLY, + 1, + "https://mcp.example.com/.well-known/oauth-protected-resource" + )] + #[case::protected_metadata_too_many_requests( + StatusCode::TOO_MANY_REQUESTS, + 1, + "https://mcp.example.com/.well-known/oauth-protected-resource" + )] + #[case::authorization_request_timeout( + StatusCode::REQUEST_TIMEOUT, + 2, + "https://auth.example.com/.well-known/oauth-authorization-server" + )] + #[case::authorization_too_early( + StatusCode::TOO_EARLY, + 2, + "https://auth.example.com/.well-known/oauth-authorization-server" + )] + #[case::authorization_too_many_requests( + StatusCode::TOO_MANY_REQUESTS, + 2, + "https://auth.example.com/.well-known/oauth-authorization-server" + )] + #[tokio::test] + async fn discovery_propagates_transient_client_errors( + #[case] status: StatusCode, + #[case] successful_response_count: usize, + #[case] expected_url: &str, + ) { + let mut responses = preregistered_discovery_responses(); + responses.insert(successful_response_count, empty_response(status.as_u16())); + + let client = RecordingOAuthHttpClient::with_responses(responses); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let error = manager.resolve_metadata().await.unwrap_err(); + + assert!( + matches!( + error, + AuthError::MetadataError(ref reason) + if reason.contains(expected_url) && reason.contains(status.as_str()) + ), + "unexpected discovery error for {status}: {error}" + ); + assert_eq!(client.requests().len(), successful_response_count + 1); + } + #[tokio::test] async fn custom_http_client_handles_protected_resource_discovery() { let challenge = oauth2::http::Response::builder() @@ -3233,13 +4446,14 @@ mod tests { http_response( 200, serde_json::json!({ - "resource": "https://mcp.example.com/mcp", + "resource": "https://mcp.example.com", "authorization_servers": ["https://auth.example.com"] }), ), http_response( 200, serde_json::json!({ + "issuer": "https://auth.example.com", "authorization_endpoint": "https://auth.example.com/authorize", "token_endpoint": "https://auth.example.com/token" }), @@ -3252,9 +4466,13 @@ mod tests { .await .unwrap(); - let metadata = manager.discover_metadata().await.unwrap(); + let metadata = manager.resolve_metadata().await.unwrap().metadata; assert_eq!(metadata.token_endpoint, "https://auth.example.com/token"); + assert_eq!( + manager.discovered_resource.read().await.as_deref(), + Some("https://mcp.example.com") + ); assert_eq!( client.requests(), vec![ @@ -3281,6 +4499,1690 @@ mod tests { ); } + #[tokio::test] + async fn refresh_token_uses_discovered_protected_resource() { + let challenge = oauth2::http::Response::builder() + .status(401) + .header( + "www-authenticate", + r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#, + ) + .body(Vec::new()) + .unwrap(); + let client = RecordingOAuthHttpClient::with_responses(vec![ + challenge, + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com", + "authorization_servers": ["https://auth.example.com"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token", + "response_types_supported": ["code"], + "code_challenge_methods_supported": ["S256"] + }), + ), + http_response( + 200, + serde_json::json!({ + "access_token": "initial-access-token", + "token_type": "bearer", + "refresh_token": "initial-refresh-token", + "expires_in": 3600 + }), + ), + http_response( + 200, + serde_json::json!({ + "access_token": "refreshed-access-token", + "token_type": "bearer", + "refresh_token": "refreshed-refresh-token", + "expires_in": 3600 + }), + ), + ]); + let mut manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + manager.set_metadata(resolution.metadata); + manager.configure_client_id("test-client-id").unwrap(); + + let authorization_url = manager.get_authorization_url(&[]).await.unwrap(); + let authorization_params: HashMap = Url::parse(&authorization_url) + .unwrap() + .query_pairs() + .into_owned() + .collect(); + let state = authorization_params.get("state").unwrap(); + + manager + .exchange_code_for_token("authorization-code", state) + .await + .unwrap(); + manager.refresh_token().await.unwrap(); + + let token_requests: Vec> = client + .requests() + .into_iter() + .filter(|request| request.uri == "https://auth.example.com/token") + .map(|request| { + url::form_urlencoded::parse(&request.body) + .into_owned() + .collect() + }) + .collect(); + + assert_eq!(token_requests.len(), 2); + assert_eq!( + token_requests[0].get("grant_type").map(String::as_str), + Some("authorization_code") + ); + assert_eq!( + token_requests[1].get("grant_type").map(String::as_str), + Some("refresh_token") + ); + assert_eq!( + [ + authorization_params.get("resource").map(String::as_str), + token_requests[0].get("resource").map(String::as_str), + token_requests[1].get("resource").map(String::as_str), + ], + [Some("https://mcp.example.com"); 3], + "authorization, code exchange, and refresh must use the discovered resource audience" + ); + } + + #[tokio::test] + async fn protected_resource_metadata_supports_authorization_server_path_insertion() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + empty_response(401), + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/", + "authorization_servers": ["https://auth.example.com/tenant1"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com/tenant1", + "authorization_endpoint": "https://auth.example.com/tenant1/authorize", + "token_endpoint": "https://auth.example.com/tenant1/token" + }), + ), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let metadata = manager.resolve_metadata().await.unwrap().metadata; + + assert_eq!( + ( + metadata.issuer.as_deref(), + metadata.authorization_endpoint.as_str(), + client + .requests() + .iter() + .map(|request| request.uri.as_str()) + .collect::>(), + ), + ( + Some("https://auth.example.com/tenant1"), + "https://auth.example.com/tenant1/authorize", + vec![ + "https://mcp.example.com/", + "https://mcp.example.com/.well-known/oauth-protected-resource", + "https://auth.example.com/.well-known/oauth-authorization-server/tenant1", + ], + ) + ); + } + + #[tokio::test] + async fn protected_resource_metadata_preserves_non_root_issuer_trailing_slash() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + empty_response(401), + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/", + "authorization_servers": ["https://auth.example.com/tenant1/"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com/tenant1/", + "authorization_endpoint": "https://auth.example.com/tenant1/authorize", + "token_endpoint": "https://auth.example.com/tenant1/token" + }), + ), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let metadata = manager.resolve_metadata().await.unwrap().metadata; + + assert_eq!( + ( + metadata.issuer.as_deref(), + client.requests().last().map(|request| request.uri.as_str()), + ), + ( + Some("https://auth.example.com/tenant1/"), + Some("https://auth.example.com/.well-known/oauth-authorization-server/tenant1"), + ) + ); + } + + #[tokio::test] + async fn authorization_metadata_rejects_mismatched_issuer() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + empty_response(401), + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/", + "authorization_servers": ["https://auth.example.com/tenant1"] + }), + ), + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/", + "authorization_servers": ["https://auth.example.com/tenant1"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://evil.example.com/tenant1", + "authorization_endpoint": "https://evil.example.com/tenant1/authorize", + "token_endpoint": "https://evil.example.com/tenant1/token" + }), + ), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/", + Arc::new(client), + ) + .await + .unwrap(); + + let error = manager.resolve_metadata().await.unwrap_err(); + + assert!( + matches!( + error, + AuthError::AuthorizationServerMismatch { + ref expected_issuer, + ref received_issuer + } if expected_issuer == "https://auth.example.com/tenant1" + && received_issuer == "https://evil.example.com/tenant1" + ), + "expected authorization server issuer mismatch, got: {error:?}" + ); + } + + #[tokio::test] + async fn authorization_metadata_accepts_oidc_path_appended_issuer() { + let discovery_url = + Url::parse("https://auth.example.com/tenant1/.well-known/openid-configuration") + .unwrap(); + let metadata = AuthorizationMetadata { + issuer: Some("https://auth.example.com/tenant1".to_string()), + authorization_endpoint: "https://auth.example.com/tenant1/authorize".to_string(), + token_endpoint: "https://auth.example.com/tenant1/token".to_string(), + ..Default::default() + }; + let manager = AuthorizationManager::new("https://mcp.example.com/") + .await + .unwrap(); + + manager + .validate_authorization_metadata_issuer(&discovery_url, None, &metadata) + .unwrap(); + } + + #[test] + fn authorization_metadata_allows_only_root_trailing_slash_equivalence() { + assert!(AuthorizationManager::issuer_identifiers_match( + "https://auth.example.com/", + "https://auth.example.com" + )); + assert!(!AuthorizationManager::issuer_identifiers_match( + "https://auth.example.com/tenant1/", + "https://auth.example.com/tenant1" + )); + } + + #[tokio::test] + async fn authorization_metadata_accepts_oidc_path_inserted_issuer() { + let discovery_url = + Url::parse("https://auth.example.com/.well-known/openid-configuration/tenant1") + .unwrap(); + let metadata = AuthorizationMetadata { + issuer: Some("https://auth.example.com/tenant1".to_string()), + authorization_endpoint: "https://auth.example.com/tenant1/authorize".to_string(), + token_endpoint: "https://auth.example.com/tenant1/token".to_string(), + ..Default::default() + }; + let manager = AuthorizationManager::new("https://mcp.example.com/") + .await + .unwrap(); + + manager + .validate_authorization_metadata_issuer(&discovery_url, None, &metadata) + .unwrap(); + } + + #[tokio::test] + async fn authorization_metadata_allows_missing_issuer_when_configured() { + let discovery_url = + Url::parse("https://auth.example.com/.well-known/openid-configuration/tenant1") + .unwrap(); + let metadata = AuthorizationMetadata { + issuer: None, + authorization_endpoint: "https://auth.example.com/tenant1/authorize".to_string(), + token_endpoint: "https://auth.example.com/tenant1/token".to_string(), + ..Default::default() + }; + + let mut manager = AuthorizationManager::new("https://mcp.example.com/") + .await + .unwrap(); + manager.set_allow_missing_issuer(true); + + manager + .validate_authorization_metadata_issuer(&discovery_url, None, &metadata) + .unwrap(); + } + + #[tokio::test] + async fn authorization_metadata_rejects_missing_issuer_by_default() { + let discovery_url = + Url::parse("https://auth.example.com/.well-known/openid-configuration/tenant1") + .unwrap(); + let metadata = AuthorizationMetadata { + issuer: None, + authorization_endpoint: "https://auth.example.com/tenant1/authorize".to_string(), + token_endpoint: "https://auth.example.com/tenant1/token".to_string(), + ..Default::default() + }; + let manager = AuthorizationManager::new("https://mcp.example.com/") + .await + .unwrap(); + + let error = manager + .validate_authorization_metadata_issuer(&discovery_url, None, &metadata) + .unwrap_err(); + + assert!( + matches!( + error, + AuthError::AuthorizationServerMissingIssuer { ref expected_issuer } + if expected_issuer == "https://auth.example.com/tenant1" + ), + "expected missing issuer error, got: {error:?}" + ); + } + + #[tokio::test] + async fn protected_resource_metadata_supports_custom_location_and_oidc_path_append() { + let challenge = oauth2::http::Response::builder() + .status(401) + .header( + "www-authenticate", + r#"Bearer resource_metadata="/custom/metadata/location.json""#, + ) + .body(Vec::new()) + .unwrap(); + let client = RecordingOAuthHttpClient::with_responses(vec![ + challenge, + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/mcp", + "authorization_servers": ["https://auth.example.com/tenant1"] + }), + ), + empty_response(404), + empty_response(404), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com/tenant1", + "authorization_endpoint": "https://auth.example.com/tenant1/authorize", + "token_endpoint": "https://auth.example.com/tenant1/token" + }), + ), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let metadata = manager.resolve_metadata().await.unwrap().metadata; + + assert_eq!( + ( + metadata.token_endpoint.as_str(), + client + .requests() + .iter() + .map(|request| request.uri.as_str()) + .collect::>(), + ), + ( + "https://auth.example.com/tenant1/token", + vec![ + "https://mcp.example.com/mcp", + "https://mcp.example.com/custom/metadata/location.json", + "https://auth.example.com/.well-known/oauth-authorization-server/tenant1", + "https://auth.example.com/.well-known/openid-configuration/tenant1", + "https://auth.example.com/tenant1/.well-known/openid-configuration", + ], + ) + ); + assert!( + client + .requests() + .iter() + .all(|request| request.method == "GET"), + "discovery must not send non-GET requests" + ); + } + + #[rstest] + #[case::not_found(StatusCode::NOT_FOUND)] + #[case::method_not_allowed(StatusCode::METHOD_NOT_ALLOWED)] + #[tokio::test] + async fn resolve_metadata_reports_legacy_fallback_when_nothing_is_discovered( + #[case] status: StatusCode, + ) { + let client = RecordingOAuthHttpClient::with_responses(vec![ + empty_response(status.as_u16()), + empty_response(status.as_u16()), + empty_response(status.as_u16()), + empty_response(status.as_u16()), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://legacy.example.com/", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.authorization_endpoint.as_str(), + resolution.metadata.token_endpoint.as_str(), + resolution.metadata.registration_endpoint.as_deref(), + client + .requests() + .iter() + .map(|request| request.uri.as_str()) + .collect::>(), + ), + ( + AuthorizationMetadataSource::LegacyEndpointFallback, + "https://legacy.example.com/authorize", + "https://legacy.example.com/token", + Some("https://legacy.example.com/register"), + vec![ + "https://legacy.example.com/", + "https://legacy.example.com/.well-known/oauth-protected-resource", + "https://legacy.example.com/.well-known/oauth-authorization-server", + "https://legacy.example.com/.well-known/openid-configuration", + ], + ) + ); + } + + #[tokio::test] + async fn resolve_metadata_reports_protected_resource_metadata() { + let challenge = oauth2::http::Response::builder() + .status(401) + .header( + "www-authenticate", + r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#, + ) + .body(Vec::new()) + .unwrap(); + let client = RecordingOAuthHttpClient::with_responses(vec![ + challenge, + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/mcp", + "authorization_servers": ["https://auth.example.com"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token" + }), + ), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.token_endpoint.as_str(), + ), + ( + AuthorizationMetadataSource::ProtectedResourceMetadata, + "https://auth.example.com/token", + ) + ); + } + + #[tokio::test] + async fn resolve_metadata_from_challenge_uses_challenge_pointer_and_scope() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/mcp", + "authorization_servers": ["https://auth.example.com"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token" + }), + ), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let resolution = manager + .resolve_metadata_from_challenge(Some( + r#"Bearer resource_metadata="https://mcp.example.com/custom/prm.json", scope="mcp:read mcp:write""#, + )) + .await + .unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.token_endpoint.as_str(), + manager.select_scopes(None, &[]), + client.requests().first().map(|request| request.uri.clone()), + ), + ( + AuthorizationMetadataSource::ProtectedResourceMetadata, + "https://auth.example.com/token", + vec!["mcp:read".to_string(), "mcp:write".to_string()], + // discovery starts at the challenge's pointer: no probing + Some("https://mcp.example.com/custom/prm.json".to_string()), + ) + ); + } + + #[tokio::test] + async fn resolve_metadata_from_challenge_falls_back_without_metadata_pointer() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + empty_response(404), + empty_response(404), + empty_response(404), + empty_response(404), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://legacy.example.com/", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager + .resolve_metadata_from_challenge(Some(r#"Bearer scope="mcp:basic""#)) + .await + .unwrap(); + + assert_eq!( + (resolution.source, manager.select_scopes(None, &[]),), + ( + AuthorizationMetadataSource::LegacyEndpointFallback, + vec!["mcp:basic".to_string()], + ) + ); + } + + #[tokio::test] + async fn resolve_metadata_reports_authorization_server_metadata() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + empty_response(404), + empty_response(404), + empty_response(404), + http_response( + 200, + serde_json::json!({ + "issuer": "https://mcp.example.com", + "authorization_endpoint": "https://mcp.example.com/oauth/authorize", + "token_endpoint": "https://mcp.example.com/oauth/token" + }), + ), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.token_endpoint.as_str(), + ), + ( + AuthorizationMetadataSource::AuthorizationServerMetadata, + "https://mcp.example.com/oauth/token", + ) + ); + } + + #[tokio::test] + async fn resolve_metadata_reaches_the_well_known_document_past_a_non_metadata_base_url() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + // the MCP endpoint answers GET with a health payload, not metadata + http_response( + 200, + serde_json::json!({"status": "healthy", "message": "MCP server is running"}), + ), + // the well-known candidate carries the real document + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/", + "authorization_servers": ["https://auth.example.com"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token" + }), + ), + ]); + let recorder = client.clone(); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.token_endpoint.as_str(), + ), + ( + AuthorizationMetadataSource::ProtectedResourceMetadata, + "https://auth.example.com/token", + ) + ); + assert!( + recorder.requests().iter().any(|request| { + request.uri == "https://mcp.example.com/.well-known/oauth-protected-resource" + }), + "the well-known candidate was never probed: {:?}", + recorder.requests() + ); + } + + #[tokio::test] + async fn resolve_metadata_ignores_a_well_known_url_that_is_not_a_metadata_document() { + let health = || { + http_response( + 200, + serde_json::json!({"status": "healthy", "message": "MCP server is running"}), + ) + }; + let client = RecordingOAuthHttpClient::with_responses(vec![ + // the MCP endpoint answers GET with a health payload, not metadata + health(), + // so does the well-known candidate + health(), + http_response( + 200, + serde_json::json!({ + "issuer": "https://mcp.example.com", + "authorization_endpoint": "https://mcp.example.com/oauth/authorize", + "token_endpoint": "https://mcp.example.com/oauth/token" + }), + ), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.token_endpoint.as_str(), + ), + ( + AuthorizationMetadataSource::AuthorizationServerMetadata, + "https://mcp.example.com/oauth/token", + ) + ); + } + + #[tokio::test] + async fn resolve_metadata_tries_the_next_well_known_candidate_past_a_non_metadata_document() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + // the MCP endpoint answers GET with a health payload, not metadata + http_response( + 200, + serde_json::json!({"status": "healthy", "message": "MCP server is running"}), + ), + // so does the first well-known candidate + http_response( + 200, + serde_json::json!({"status": "healthy", "message": "MCP server is running"}), + ), + // the second candidate carries the real document + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/mcp", + "authorization_servers": ["https://auth.example.com"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token" + }), + ), + ]); + let recorder = client.clone(); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.token_endpoint.as_str(), + ), + ( + AuthorizationMetadataSource::ProtectedResourceMetadata, + "https://auth.example.com/token", + ) + ); + let document_requests = recorder + .requests() + .iter() + .filter(|request| { + request.uri == "https://mcp.example.com/mcp/.well-known/oauth-protected-resource" + }) + .count(); + assert_eq!( + document_requests, + 1, + "the candidate holding the document should be requested exactly once: {:?}", + recorder.requests() + ); + } + + #[tokio::test] + async fn resolve_metadata_from_challenge_reports_an_advertised_url_without_metadata() { + let mut responses = vec![http_response(200, serde_json::json!({}))]; + // enough responses for the fallback to reach the legacy endpoints, so that + // treating the document as a soft failure would resolve rather than error + responses.extend(std::iter::repeat_with(|| empty_response(404)).take(8)); + let client = RecordingOAuthHttpClient::with_responses(responses); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client), + ) + .await + .unwrap(); + + let error = manager + .resolve_metadata_from_challenge(Some( + r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#, + )) + .await + .unwrap_err(); + + assert_eq!( + error.to_string(), + "Metadata error: the server advertised https://mcp.example.com/.well-known/oauth-protected-resource as protected resource metadata, but the document carries neither `resource` nor an authorization server reference" + ); + } + + #[tokio::test] + async fn resolve_metadata_tries_the_next_well_known_candidate_past_an_unusable_document() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + // the MCP endpoint answers GET with a health payload, not metadata + http_response( + 200, + serde_json::json!({"status": "healthy", "message": "MCP server is running"}), + ), + // a catch-all handler answers the first candidate with its own error + // shape, which carries a `resource` that only validation rejects + http_response( + 200, + serde_json::json!({ + "error": "not_found", + "resource": "/.well-known/oauth-protected-resource" + }), + ), + // the second candidate carries the real document + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/mcp", + "authorization_servers": ["https://auth.example.com"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token" + }), + ), + ]); + let recorder = client.clone(); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.token_endpoint.as_str(), + ), + ( + AuthorizationMetadataSource::ProtectedResourceMetadata, + "https://auth.example.com/token", + ) + ); + assert!( + recorder.requests().iter().any(|request| { + request.uri == "https://mcp.example.com/mcp/.well-known/oauth-protected-resource" + }), + "the candidate after the rejected one was never probed: {:?}", + recorder.requests() + ); + } + + #[tokio::test] + async fn resolve_metadata_probes_the_candidates_past_an_advertised_url_that_is_not_served() { + let challenge = oauth2::http::Response::builder() + .status(401) + .header( + "www-authenticate", + r#"Bearer resource_metadata="https://mcp.example.com/prm""#, + ) + .body(Vec::new()) + .unwrap(); + let client = RecordingOAuthHttpClient::with_responses(vec![ + challenge, + // the advertised url is not where the document is served + empty_response(404), + // neither is the first candidate + empty_response(404), + // the second candidate carries the document + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/mcp", + "authorization_servers": ["https://auth.example.com"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token" + }), + ), + ]); + let recorder = client.clone(); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.token_endpoint.as_str(), + ), + ( + AuthorizationMetadataSource::ProtectedResourceMetadata, + "https://auth.example.com/token", + ) + ); + assert!( + recorder.requests().iter().any(|request| { + request.uri == "https://mcp.example.com/mcp/.well-known/oauth-protected-resource" + }), + "the candidates were skipped after the advertised url answered 404: {:?}", + recorder.requests() + ); + } + + #[tokio::test] + async fn resolve_metadata_requests_an_advertised_url_that_is_also_a_candidate_once() { + let challenge = oauth2::http::Response::builder() + .status(401) + .header( + "www-authenticate", + r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#, + ) + .body(Vec::new()) + .unwrap(); + let mut responses = vec![ + // the MCP endpoint answers GET with a health payload, not metadata + http_response( + 200, + serde_json::json!({"status": "healthy", "message": "MCP server is running"}), + ), + // the first candidate points at the last candidate of the same run + challenge, + ]; + // the document is served nowhere, so the run walks every candidate and + // settles on the legacy endpoints + responses.extend(std::iter::repeat_with(|| empty_response(404)).take(10)); + let client = RecordingOAuthHttpClient::with_responses(responses); + let recorder = client.clone(); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + resolution.source, + AuthorizationMetadataSource::LegacyEndpointFallback + ); + let advertised_requests = recorder + .requests() + .iter() + .filter(|request| { + request.uri == "https://mcp.example.com/.well-known/oauth-protected-resource" + }) + .count(); + assert_eq!( + advertised_requests, + 1, + "the url the challenge named was requested again as a candidate: {:?}", + recorder.requests() + ); + } + + #[tokio::test] + async fn resolve_metadata_requests_an_authorization_server_named_twice_once() { + let mut responses = vec![ + // the MCP endpoint answers GET with a health payload, not metadata + http_response( + 200, + serde_json::json!({"status": "healthy", "message": "MCP server is running"}), + ), + // the document names the same authorization server in both the singular + // draft field and the list + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/", + "authorization_server": "https://auth.example.com", + "authorization_servers": ["https://auth.example.com"] + }), + ), + ]; + // the authorization server publishes no metadata; more 404s than the walk + // can use, so that requesting it twice would show up in the sequence + responses.extend(std::iter::repeat_with(|| empty_response(404)).take(10)); + let client = RecordingOAuthHttpClient::with_responses(responses); + let recorder = client.clone(); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/", + Arc::new(client), + ) + .await + .unwrap(); + + manager.resolve_metadata().await.unwrap_err(); + + let requests = recorder.requests(); + let requested: Vec<_> = requests + .iter() + .map(|request| request.uri.as_str()) + .collect(); + assert_eq!( + requested, + [ + "https://mcp.example.com/", + "https://mcp.example.com/.well-known/oauth-protected-resource", + "https://auth.example.com/.well-known/oauth-authorization-server", + "https://auth.example.com/.well-known/openid-configuration", + ] + ); + } + + #[tokio::test] + async fn resolve_metadata_reports_named_authorization_servers_that_publish_no_metadata() { + let mut responses = vec![ + // the MCP endpoint answers GET with a health payload, not metadata + http_response( + 200, + serde_json::json!({"status": "healthy", "message": "MCP server is running"}), + ), + // the document names two authorization servers + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/", + "authorization_servers": ["https://auth.example.com", "https://sso.example.com"] + }), + ), + ]; + // neither publishes metadata; the surplus 404s would let the run reach the + // resource host's own well-known urls and the legacy endpoints if it kept going + responses.extend(std::iter::repeat_with(|| empty_response(404)).take(10)); + let client = RecordingOAuthHttpClient::with_responses(responses); + let recorder = client.clone(); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/", + Arc::new(client), + ) + .await + .unwrap(); + + let error = manager.resolve_metadata().await.unwrap_err(); + + assert_eq!( + error.to_string(), + "Metadata error: protected resource metadata at https://mcp.example.com/.well-known/oauth-protected-resource names authorization servers https://auth.example.com, https://sso.example.com, but none published usable metadata" + ); + let requests = recorder.requests(); + let requested: Vec<_> = requests + .iter() + .map(|request| request.uri.as_str()) + .collect(); + assert_eq!( + requested, + [ + "https://mcp.example.com/", + "https://mcp.example.com/.well-known/oauth-protected-resource", + "https://auth.example.com/.well-known/oauth-authorization-server", + "https://auth.example.com/.well-known/openid-configuration", + "https://sso.example.com/.well-known/oauth-authorization-server", + "https://sso.example.com/.well-known/openid-configuration", + ] + ); + } + + #[tokio::test] + async fn resolve_metadata_keeps_falling_back_when_the_document_names_no_authorization_server() { + let mut responses = vec![ + // the MCP endpoint answers GET with a health payload, not metadata + http_response( + 200, + serde_json::json!({"status": "healthy", "message": "MCP server is running"}), + ), + // the document says nothing about where to authorize + http_response( + 200, + serde_json::json!({"resource": "https://mcp.example.com/"}), + ), + ]; + responses.extend(std::iter::repeat_with(|| empty_response(404)).take(10)); + let client = RecordingOAuthHttpClient::with_responses(responses); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + resolution.source, + AuthorizationMetadataSource::LegacyEndpointFallback + ); + } + + #[tokio::test] + async fn resolve_metadata_tries_the_next_candidate_past_a_pointer_found_while_guessing() { + let challenge = oauth2::http::Response::builder() + .status(401) + .header( + "www-authenticate", + r#"Bearer resource_metadata="https://mcp.example.com/prm""#, + ) + .body(Vec::new()) + .unwrap(); + let client = RecordingOAuthHttpClient::with_responses(vec![ + // the MCP endpoint answers GET with a health payload, not metadata + http_response( + 200, + serde_json::json!({"status": "healthy", "message": "MCP server is running"}), + ), + // auth sits in front of the first candidate and its challenge names a url + challenge, + // which serves something that is not the document + http_response(200, serde_json::json!({})), + // the second candidate carries the real document + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/mcp", + "authorization_servers": ["https://auth.example.com"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token" + }), + ), + ]); + let recorder = client.clone(); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager.resolve_metadata().await.unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.token_endpoint.as_str(), + ), + ( + AuthorizationMetadataSource::ProtectedResourceMetadata, + "https://auth.example.com/token", + ) + ); + let requests = recorder.requests(); + let requested: Vec<_> = requests + .iter() + .map(|request| request.uri.as_str()) + .collect(); + assert_eq!( + requested, + [ + "https://mcp.example.com/mcp", + "https://mcp.example.com/.well-known/oauth-protected-resource/mcp", + "https://mcp.example.com/prm", + "https://mcp.example.com/mcp/.well-known/oauth-protected-resource", + "https://auth.example.com/.well-known/oauth-authorization-server", + ] + ); + } + + #[rstest] + #[case::protected_resource_metadata( + AuthorizationMetadataSource::ProtectedResourceMetadata, + true + )] + #[case::authorization_server_metadata( + AuthorizationMetadataSource::AuthorizationServerMetadata, + true + )] + #[case::legacy_endpoint_fallback(AuthorizationMetadataSource::LegacyEndpointFallback, false)] + fn is_discovered_is_false_only_for_the_legacy_fallback( + #[case] source: AuthorizationMetadataSource, + #[case] expected: bool, + ) { + assert_eq!(source.is_discovered(), expected); + } + + fn preregistered_as_metadata_response() -> HttpResponse { + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token", + "registration_endpoint": "https://auth.example.com/register", + "scopes_supported": ["read", "write", "offline_access"] + }), + ) + } + + /// discovery responses for the preregistered-client tests: a 401 challenge + /// pointing at protected resource metadata, the PRM document, then the + /// authorization server metadata. + fn preregistered_discovery_responses() -> Vec { + let challenge = oauth2::http::Response::builder() + .status(401) + .header( + "www-authenticate", + r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#, + ) + .body(Vec::new()) + .unwrap(); + vec![ + challenge, + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/mcp", + "authorization_servers": ["https://auth.example.com"] + }), + ), + preregistered_as_metadata_response(), + ] + } + + fn auth_url_query(auth_url: &str) -> HashMap { + Url::parse(auth_url) + .unwrap() + .query_pairs() + .map(|(k, v)| (k.into_owned(), v.into_owned())) + .collect() + } + + #[tokio::test] + async fn preregistered_client_skips_registration_endpoint() { + let client = RecordingOAuthHttpClient::with_responses(preregistered_discovery_responses()); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_preregistered_client("preregistered-client") + .with_client_secret("secret") + .with_scopes(["read"]); + state.start_authorization(request).await.unwrap(); + + // the registration endpoint was advertised but must not be called + let requests = client.requests(); + assert!( + requests + .iter() + .all(|request| !request.uri.contains("/register")), + "registration endpoint should not be called: {requests:?}" + ); + + let auth_url = state.get_authorization_url().await.unwrap(); + let query = auth_url_query(&auth_url); + assert_eq!(query.get("client_id").unwrap(), "preregistered-client"); + assert!(matches!(state, super::OAuthState::Session(_))); + } + + #[tokio::test] + async fn preregistered_client_uses_metadata_resource_when_server_url_has_query() { + let client = RecordingOAuthHttpClient::with_responses(preregistered_discovery_responses()); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp?oauth=initialize", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_preregistered_client("preregistered-client"); + state.start_authorization(request).await.unwrap(); + + let auth_url = state.get_authorization_url().await.unwrap(); + let query = auth_url_query(&auth_url); + assert_eq!( + query.get("resource").map(String::as_str), + Some("https://mcp.example.com/mcp") + ); + assert_eq!( + client.requests()[0].uri, + "https://mcp.example.com/mcp?oauth=initialize" + ); + } + + #[tokio::test] + async fn authorization_session_selects_default_scopes_when_none_provided() { + let client = RecordingOAuthHttpClient::with_responses(preregistered_discovery_responses()); + let mut manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client), + ) + .await + .unwrap(); + manager.metadata = Some(manager.resolve_metadata().await.unwrap().metadata); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_preregistered_client("preregistered-client"); + let session = match AuthorizationSession::new(manager, request).await { + Ok(session) => session, + Err((_, error)) => panic!("authorization session creation failed: {error}"), + }; + + // Empty request scopes fall back to the discovered scopes_supported. + let query = auth_url_query(&session.auth_url); + assert_eq!(query.get("scope").unwrap(), "read write offline_access"); + } + + #[tokio::test] + async fn preregistered_client_uses_explicit_scopes_and_adds_offline_access() { + let client = RecordingOAuthHttpClient::with_responses(preregistered_discovery_responses()); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_preregistered_client("preregistered-client") + .with_scopes(["read"]); + state.start_authorization(request).await.unwrap(); + + // explicit scopes are preserved; offline_access is appended per SEP-2207 + let auth_url = state.get_authorization_url().await.unwrap(); + let query = auth_url_query(&auth_url); + assert_eq!(query.get("scope").unwrap(), "read offline_access"); + } + + #[tokio::test] + async fn preregistered_client_recovers_unauthorized_state_after_discovery_failure() { + // first discovery attempt fails: protected resource metadata reports a + // mismatched resource identifier, which discover_metadata rejects + let challenge = oauth2::http::Response::builder() + .status(401) + .header( + "www-authenticate", + r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#, + ) + .body(Vec::new()) + .unwrap(); + let client = RecordingOAuthHttpClient::with_responses(vec![ + challenge, + http_response( + 200, + serde_json::json!({ + "resource": "https://other.example.com/mcp", + "authorization_servers": ["https://auth.example.com"] + }), + ), + ]); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_preregistered_client("preregistered-client") + .with_scopes(["read"]); + let err = state + .start_authorization(request.clone()) + .await + .unwrap_err(); + assert!(!matches!(err, AuthError::InternalError(_)), "{err:?}"); + assert!( + matches!(state, super::OAuthState::Unauthorized(_)), + "state should return to Unauthorized after a transient failure" + ); + + // retrying with the same state succeeds once the server responds + client + .responses + .lock() + .unwrap() + .extend(preregistered_discovery_responses()); + state.start_authorization(request).await.unwrap(); + assert!(matches!(state, super::OAuthState::Session(_))); + } + + fn cimd_as_metadata_response() -> HttpResponse { + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token", + "registration_endpoint": "https://auth.example.com/register", + "scopes_supported": ["read", "write", "offline_access"], + "client_id_metadata_document_supported": true + }), + ) + } + + /// discovery responses like [`preregistered_discovery_responses`] but the + /// authorization server advertises CIMD support. + fn cimd_discovery_responses() -> Vec { + let challenge = oauth2::http::Response::builder() + .status(401) + .header( + "www-authenticate", + r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#, + ) + .body(Vec::new()) + .unwrap(); + vec![ + challenge, + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/mcp", + "authorization_servers": ["https://auth.example.com"] + }), + ), + cimd_as_metadata_response(), + ] + } + + #[tokio::test] + async fn preregistered_client_takes_priority_over_cimd() { + // server supports CIMD and the request carries both pre-registered + // credentials and a client metadata URL: pre-registration wins + let client = RecordingOAuthHttpClient::with_responses(cimd_discovery_responses()); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_preregistered_client("preregistered-client") + .with_client_metadata_url("https://client.example.com/client-metadata.json") + .with_scopes(["read"]); + state.start_authorization(request).await.unwrap(); + + let requests = client.requests(); + assert!( + requests + .iter() + .all(|request| !request.uri.contains("/register")), + "registration endpoint should not be called: {requests:?}" + ); + let auth_url = state.get_authorization_url().await.unwrap(); + let query = auth_url_query(&auth_url); + assert_eq!(query.get("client_id").unwrap(), "preregistered-client"); + } + + #[tokio::test] + async fn cimd_used_when_server_advertises_support() { + let client = RecordingOAuthHttpClient::with_responses(cimd_discovery_responses()); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_client_metadata_url("https://client.example.com/client-metadata.json") + .with_scopes(["read"]); + state.start_authorization(request).await.unwrap(); + + // CIMD takes priority over the advertised registration endpoint + let requests = client.requests(); + assert!( + requests + .iter() + .all(|request| !request.uri.contains("/register")), + "registration endpoint should not be called: {requests:?}" + ); + let auth_url = state.get_authorization_url().await.unwrap(); + let query = auth_url_query(&auth_url); + assert_eq!( + query.get("client_id").unwrap(), + "https://client.example.com/client-metadata.json" + ); + } + + #[tokio::test] + async fn cimd_falls_back_to_dcr_when_server_lacks_support() { + // server does not advertise client_id_metadata_document_supported, so + // the client metadata URL is ignored and DCR is used instead + let mut responses = preregistered_discovery_responses(); + responses.push(http_response( + 201, + serde_json::json!({ + "client_id": "dcr-client", + "redirect_uris": ["http://localhost:8080/callback"] + }), + )); + let client = RecordingOAuthHttpClient::with_responses(responses); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_client_metadata_url("https://client.example.com/client-metadata.json") + .with_scopes(["read"]); + state.start_authorization(request).await.unwrap(); + + let requests = client.requests(); + assert!( + requests + .iter() + .any(|request| request.uri.contains("/register")), + "registration endpoint should be called: {requests:?}" + ); + let auth_url = state.get_authorization_url().await.unwrap(); + let query = auth_url_query(&auth_url); + assert_eq!(query.get("client_id").unwrap(), "dcr-client"); + } + + #[tokio::test] + async fn dcr_used_when_no_identity_material_is_provided() { + let mut responses = preregistered_discovery_responses(); + responses.push(http_response( + 201, + serde_json::json!({ + "client_id": "dcr-client", + "redirect_uris": ["http://localhost:8080/callback"] + }), + )); + let client = RecordingOAuthHttpClient::with_responses(responses); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_client_name("test-client") + .with_scopes(["read"]); + state.start_authorization(request).await.unwrap(); + + let auth_url = state.get_authorization_url().await.unwrap(); + let query = auth_url_query(&auth_url); + assert_eq!(query.get("client_id").unwrap(), "dcr-client"); + assert!(matches!(state, super::OAuthState::Session(_))); + } + + #[tokio::test] + async fn dcr_registration_uses_requested_application_type() { + let mut responses = preregistered_discovery_responses(); + responses.push(http_response( + 201, + serde_json::json!({ + "client_id": "dcr-client", + "redirect_uris": ["http://localhost:8080/callback"] + }), + )); + let client = RecordingOAuthHttpClient::with_responses(responses); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_client_name("test-client") + .with_application_type("web") + .with_scopes(["read"]); + state.start_authorization(request).await.unwrap(); + + // SEP-837: the requested application_type must be sent in the DCR request + let requests = client.requests(); + let registration = requests + .iter() + .find(|request| request.uri.contains("/register")) + .expect("registration endpoint should be called"); + let body: serde_json::Value = serde_json::from_slice(®istration.body).unwrap(); + assert_eq!(body.get("application_type").unwrap(), "web"); + } + + #[tokio::test] + async fn cimd_rejects_non_https_client_metadata_url_and_recovers_state() { + let client = RecordingOAuthHttpClient::with_responses(cimd_discovery_responses()); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_client_metadata_url("http://client.example.com/client-metadata.json") + .with_scopes(["read"]); + let err = state.start_authorization(request).await.unwrap_err(); + assert!(matches!(err, AuthError::RegistrationFailed(_)), "{err:?}"); + assert!( + matches!(state, super::OAuthState::Unauthorized(_)), + "state should return to Unauthorized after a registration failure" + ); + } + + #[tokio::test] + async fn dcr_recovers_unauthorized_state_after_registration_failure() { + // discovery succeeds, but the registration endpoint rejects the request + let mut responses = preregistered_discovery_responses(); + responses.push(http_response( + 400, + serde_json::json!({"error": "invalid_client_metadata"}), + )); + let client = RecordingOAuthHttpClient::with_responses(responses); + let mut state = super::OAuthState::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let request = AuthorizationRequest::new("http://localhost:8080/callback") + .with_client_name("test-client") + .with_scopes(["read"]); + let err = state + .start_authorization(request.clone()) + .await + .unwrap_err(); + assert!(matches!(err, AuthError::RegistrationFailed(_)), "{err:?}"); + assert!( + matches!(state, super::OAuthState::Unauthorized(_)), + "state should return to Unauthorized after a registration failure" + ); + + // retrying with the same state succeeds once the server accepts + // registration (discovery runs again on retry) + { + let mut responses = client.responses.lock().unwrap(); + responses.extend(preregistered_discovery_responses()); + responses.push_back(http_response( + 201, + serde_json::json!({ + "client_id": "dcr-client", + "redirect_uris": ["http://localhost:8080/callback"] + }), + )); + } + state.start_authorization(request).await.unwrap(); + assert!(matches!(state, super::OAuthState::Session(_))); + } + #[tokio::test] async fn discovery_get_follows_same_origin_redirects() { let client = RecordingOAuthHttpClient::with_responses(vec![ @@ -3362,6 +6264,7 @@ mod tests { "resource": "https://mcp.example.com/mcp", "authorization_servers": [ "http://169.254.169.254/latest/meta-data/", + "http://127.0.0.1:8080/tenant1", "https://auth.example.com" ] }), @@ -3369,6 +6272,7 @@ mod tests { http_response( 200, serde_json::json!({ + "issuer": "https://auth.example.com", "authorization_endpoint": "https://auth.example.com/authorize", "token_endpoint": "https://auth.example.com/token" }), @@ -3381,7 +6285,7 @@ mod tests { .await .unwrap(); - let metadata = manager.discover_metadata().await.unwrap(); + let metadata = manager.resolve_metadata().await.unwrap().metadata; let requests = client.requests(); assert_eq!( @@ -3403,6 +6307,63 @@ mod tests { ); } + #[tokio::test] + async fn allows_loopback_authorization_server_when_resource_is_loopback() { + let challenge = oauth2::http::Response::builder() + .status(401) + .header( + "www-authenticate", + r#"Bearer resource_metadata="http://localhost/custom-metadata.json""#, + ) + .body(Vec::new()) + .unwrap(); + let client = RecordingOAuthHttpClient::with_responses(vec![ + challenge, + http_response( + 200, + serde_json::json!({ + "resource": "http://localhost/mcp", + "authorization_servers": ["http://127.0.0.1:8080/tenant1"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "http://127.0.0.1:8080/tenant1", + "authorization_endpoint": "http://127.0.0.1:8080/tenant1/authorize", + "token_endpoint": "http://127.0.0.1:8080/tenant1/token" + }), + ), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "http://localhost/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let metadata = manager.resolve_metadata().await.unwrap().metadata; + + assert_eq!( + ( + metadata.issuer.as_deref(), + client + .requests() + .iter() + .map(|request| request.uri.as_str()) + .collect::>(), + ), + ( + Some("http://127.0.0.1:8080/tenant1"), + vec![ + "http://localhost/mcp", + "http://localhost/custom-metadata.json", + "http://127.0.0.1:8080/.well-known/oauth-authorization-server/tenant1", + ], + ) + ); + } + #[tokio::test] async fn protected_resource_discovery_rejects_mismatched_resource() { let challenge = oauth2::http::Response::builder() @@ -3430,7 +6391,7 @@ mod tests { .await .unwrap(); - let error = manager.discover_metadata().await.unwrap_err(); + let error = manager.resolve_metadata().await.unwrap_err(); assert!( matches!(error, AuthError::MetadataError(ref message) if message.contains("resource mismatch")), @@ -3465,7 +6426,7 @@ mod tests { .await .unwrap(); - let error = manager.discover_metadata().await.unwrap_err(); + let error = manager.resolve_metadata().await.unwrap_err(); assert!( matches!(error, AuthError::MetadataError(ref message) if message.contains("missing required resource")), @@ -3475,23 +6436,71 @@ mod tests { } #[test] - fn resource_identifier_matching_allows_only_root_trailing_slash_difference() { - assert!(AuthorizationManager::resource_identifiers_match( - "https://mcp.example.com/", - "https://mcp.example.com" + fn resource_identifier_matching_allows_matching_host_or_parent_path() { + assert!(AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/").unwrap(), + &Url::parse("https://mcp.example.com").unwrap() )); - assert!(AuthorizationManager::resource_identifiers_match( - "https://mcp.example.com", - "https://mcp.example.com/" + assert!(AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com").unwrap(), + &Url::parse("https://mcp.example.com/").unwrap() + )); + assert!(AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp").unwrap(), + &Url::parse("https://mcp.example.com").unwrap() + )); + assert!(AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp/tools").unwrap(), + &Url::parse("https://mcp.example.com/mcp").unwrap() + )); + assert!(AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp?query=param").unwrap(), + &Url::parse("https://mcp.example.com/mcp?query=param").unwrap() + )); + assert!(AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp?oauth=initialize").unwrap(), + &Url::parse("https://mcp.example.com/mcp").unwrap() + )); + assert!(AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp").unwrap(), + &Url::parse("https://mcp.example.com/mcp?query=value1").unwrap() + )); + assert!(AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp?query=value1").unwrap(), + &Url::parse("https://mcp.example.com/mcp?query=value2").unwrap() + )); + assert!(AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp/tools?oauth=initialize").unwrap(), + &Url::parse("https://mcp.example.com/mcp?query=value1").unwrap() )); - assert!(!AuthorizationManager::resource_identifiers_match( - "https://mcp.example.com/mcp", - "https://mcp.example.com/mcp/" + assert!(!AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp").unwrap(), + &Url::parse("https://mcp.example.com/mcp/").unwrap() )); - assert!(!AuthorizationManager::resource_identifiers_match( - "https://mcp.example.com/mcp", - "https://real.example.com/mcp" + assert!(!AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp-tools").unwrap(), + &Url::parse("https://mcp.example.com/mcp").unwrap() + )); + assert!(!AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp").unwrap(), + &Url::parse("https://mcp.example.com/mcp-tools").unwrap() + )); + assert!(!AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp").unwrap(), + &Url::parse("https://mcp.example.com/mcp/tools").unwrap() + )); + assert!(!AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp").unwrap(), + &Url::parse("https://real.example.com/mcp").unwrap() + )); + assert!(!AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp-tools?oauth=initialize").unwrap(), + &Url::parse("https://mcp.example.com/mcp?query=value1").unwrap() + )); + assert!(!AuthorizationManager::is_resource_identifier_valid( + &Url::parse("https://mcp.example.com/mcp?oauth=initialize").unwrap(), + &Url::parse("https://real.example.com/mcp?oauth=initialize").unwrap() )); } @@ -3911,17 +6920,27 @@ mod tests { fn test_stored_authorization_state_serialization() { let pkce = PkceCodeVerifier::new("my-verifier".to_string()); let csrf = CsrfToken::new("my-csrf".to_string()); - let state = StoredAuthorizationState::new(&pkce, &csrf); + let state = StoredAuthorizationState::new(&pkce, &csrf) + .with_requested_scopes(vec!["read".to_string(), "write".to_string()]); let json = serde_json::to_string(&state).unwrap(); let deserialized: StoredAuthorizationState = serde_json::from_str(&json).unwrap(); assert_eq!(deserialized.pkce_verifier, "my-verifier"); assert_eq!(deserialized.csrf_token, "my-csrf"); + assert_eq!(deserialized.requested_scopes, vec!["read", "write"]); assert_eq!(deserialized.expected_issuer, None); assert!(!deserialized.require_issuer); } + #[test] + fn stored_authorization_state_defaults_requested_scopes_when_absent() { + let json = r#"{"pkce_verifier":"v","csrf_token":"c","created_at":1}"#; + let state: StoredAuthorizationState = serde_json::from_str(json).unwrap(); + + assert!(state.requested_scopes.is_empty()); + } + #[test] fn test_stored_authorization_state_records_expected_issuer() { let pkce = PkceCodeVerifier::new("my-verifier".to_string()); @@ -3973,6 +6992,7 @@ mod tests { token_response: Some(token_response), granted_scopes: vec![], token_received_at: None, + issuer: None, }; let debug_output = format!("{:?}", creds); @@ -4158,6 +7178,34 @@ mod tests { mgr } + #[tokio::test] + async fn initialize_from_store_clears_dcr_credentials_when_issuer_changes() { + let store = InMemoryCredentialStore::new(); + store + .save(StoredCredentials { + client_id: "dcr-client".to_string(), + token_response: Some(make_token_response("old-token", Some(3600))), + granted_scopes: vec![], + token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: Some("https://old.example.com".to_string()), + }) + .await + .unwrap(); + let mut manager = manager_with_metadata(Some(AuthorizationMetadata { + authorization_endpoint: "https://new.example.com/authorize".to_string(), + token_endpoint: "https://new.example.com/token".to_string(), + issuer: Some("https://new.example.com".to_string()), + ..Default::default() + })) + .await; + manager.set_credential_store(store.clone()); + + let initialized = manager.initialize_from_store().await.unwrap(); + let credentials_cleared = store.load().await.unwrap().is_none(); + + assert_eq!((initialized, credentials_cleared), (false, true)); + } + fn test_client_config() -> OAuthClientConfig { OAuthClientConfig { client_id: "my-client".to_string(), @@ -4304,19 +7352,43 @@ mod tests { assert!(manager.validate_server_metadata("code").is_err()); } - #[tokio::test] - async fn test_validate_as_metadata_passes_without_pkce_s256() { - let mut manager = AuthorizationManager::new("https://example.com") - .await - .unwrap(); - let metadata = AuthorizationMetadata { + fn as_metadata_with_pkce(methods: Option>) -> AuthorizationMetadata { + AuthorizationMetadata { authorization_endpoint: "https://auth.example.com/authorize".to_string(), token_endpoint: "https://auth.example.com/token".to_string(), response_types_supported: Some(vec!["code".to_string()]), - code_challenge_methods_supported: Some(vec!["plain".to_string()]), + code_challenge_methods_supported: methods, ..Default::default() - }; - manager.set_metadata(metadata); + } + } + + #[tokio::test] + async fn test_validate_as_metadata_rejects_without_pkce_s256() { + let mut manager = AuthorizationManager::new("https://example.com") + .await + .unwrap(); + manager.set_metadata(as_metadata_with_pkce(Some(vec!["plain".to_string()]))); + assert!(matches!( + manager.validate_server_metadata("code"), + Err(AuthError::PkceUnsupported) + )); + } + + #[tokio::test] + async fn test_validate_as_metadata_allows_absent_pkce_methods_by_default() { + let mut manager = AuthorizationManager::new("https://example.com") + .await + .unwrap(); + manager.set_metadata(as_metadata_with_pkce(None)); + assert!(manager.validate_server_metadata("code").is_ok()); + } + + #[tokio::test] + async fn test_validate_as_metadata_passes_with_pkce_s256() { + let mut manager = AuthorizationManager::new("https://example.com") + .await + .unwrap(); + manager.set_metadata(as_metadata_with_pkce(Some(vec!["S256".to_string()]))); assert!(manager.validate_server_metadata("code").is_ok()); } @@ -4388,6 +7460,64 @@ mod tests { assert!(scope.contains("write")); } + #[tokio::test] + async fn authorization_url_uses_discovered_resource() { + let base_url = "https://mcp.example.com/mcp"; + let auth_endpoint = "https://auth.example.com/authorize"; + let mut manager = AuthorizationManager::new(base_url).await.unwrap(); + + let metadata = AuthorizationMetadata { + authorization_endpoint: auth_endpoint.to_string(), + token_endpoint: "https://auth.example.com/token".to_string(), + registration_endpoint: None, + issuer: None, + jwks_uri: None, + scopes_supported: None, + response_types_supported: Some(vec!["code".to_string()]), + code_challenge_methods_supported: Some(vec!["S256".to_string()]), + additional_fields: std::collections::HashMap::new(), + }; + manager.set_metadata(metadata); + manager.configure_client_id("test-client-id").unwrap(); + *manager.discovered_resource.write().await = Some("https://mcp.example.com".to_string()); + + let auth_url = manager.get_authorization_url(&["read"]).await.unwrap(); + let parsed = Url::parse(&auth_url).unwrap(); + let params: std::collections::HashMap<_, _> = parsed.query_pairs().collect(); + + assert_eq!( + params.get("resource").map(|v| v.as_ref()), + Some("https://mcp.example.com") + ); + } + + #[tokio::test] + async fn authorization_url_uses_default_resource_without_protected_resource_document() { + let base_url = "https://mcp.example.com/mcp"; + let auth_endpoint = "https://auth.example.com/authorize"; + let mut manager = AuthorizationManager::new(base_url).await.unwrap(); + + let metadata = AuthorizationMetadata { + authorization_endpoint: auth_endpoint.to_string(), + token_endpoint: "https://auth.example.com/token".to_string(), + registration_endpoint: None, + issuer: None, + jwks_uri: None, + scopes_supported: None, + response_types_supported: Some(vec!["code".to_string()]), + code_challenge_methods_supported: Some(vec!["S256".to_string()]), + additional_fields: std::collections::HashMap::new(), + }; + manager.set_metadata(metadata); + manager.configure_client_id("test-client-id").unwrap(); + + let auth_url = manager.get_authorization_url(&["read"]).await.unwrap(); + let parsed = Url::parse(&auth_url).unwrap(); + let params: std::collections::HashMap<_, _> = parsed.query_pairs().collect(); + + assert_eq!(params.get("resource").map(|v| v.as_ref()), Some(base_url)); + } + #[test] fn authorization_callback_parses_optional_issuer() { let callback = AuthorizationCallback::from_redirect_url( @@ -4558,6 +7688,75 @@ mod tests { ); } + #[rstest] + #[case::missing_required_issuer(None)] + #[case::mismatched_issuer(Some("https://evil.example.com"))] + #[tokio::test] + async fn invalid_issuer_does_not_consume_authorization_state( + #[case] invalid_issuer: Option<&str>, + ) { + let client = RecordingOAuthHttpClient::with_responses(vec![http_response( + 200, + serde_json::json!({ + "access_token": "access-token", + "token_type": "bearer", + "expires_in": 3600 + }), + )]); + let mut manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + manager.set_metadata(AuthorizationMetadata { + authorization_endpoint: "https://auth.example.com/authorize".to_string(), + token_endpoint: "https://auth.example.com/token".to_string(), + issuer: Some("https://auth.example.com".to_string()), + ..Default::default() + }); + manager.configure_client_id("test-client-id").unwrap(); + + let pkce = PkceCodeVerifier::new("verifier".to_string()); + let csrf = CsrfToken::new("csrf".to_string()); + let state = StoredAuthorizationState::new_with_expected_issuer( + &pkce, + &csrf, + Some("https://auth.example.com".to_string()), + true, + ); + manager.state_store.save("csrf", state).await.unwrap(); + + manager + .exchange_code_for_token_with_issuer("forged-code", "csrf", invalid_issuer) + .await + .unwrap_err(); + + assert!( + manager.state_store.load("csrf").await.unwrap().is_some(), + "issuer validation failures must leave state available for the legitimate callback" + ); + assert!( + client.requests().is_empty(), + "issuer validation failures must not reach the token endpoint" + ); + + manager + .exchange_code_for_token_with_issuer( + "legitimate-code", + "csrf", + Some("https://auth.example.com"), + ) + .await + .unwrap(); + + assert!( + manager.state_store.load("csrf").await.unwrap().is_none(), + "a valid callback must consume its one-time authorization state" + ); + assert_eq!(client.requests().len(), 1); + } + // -- scope management -- #[test] @@ -4695,6 +7894,104 @@ mod tests { assert!(scopes.contains(&"email".to_string())); } + // -- SEP-2350: client-side scope accumulation in step-up authorization -- + + #[tokio::test] + async fn select_scopes_unions_challenge_with_previously_requested() { + let mgr = manager_with_metadata(None).await; + *mgr.current_scopes.write().await = vec!["read".to_string()]; + + let scopes = mgr.select_scopes(Some("write"), &[]); + + assert_eq!(scopes, vec!["read".to_string(), "write".to_string()]); + } + + #[tokio::test] + async fn select_scopes_does_not_replace_previously_requested_with_challenge() { + let mgr = manager_with_metadata(None).await; + *mgr.current_scopes.write().await = vec!["read".to_string(), "profile".to_string()]; + + let scopes = mgr.select_scopes(Some("write"), &[]); + + assert!(scopes.contains(&"read".to_string())); + assert!(scopes.contains(&"profile".to_string())); + assert!(scopes.contains(&"write".to_string())); + } + + #[tokio::test] + async fn select_scopes_accumulates_across_multiple_step_up_rounds() { + let mgr = manager_with_metadata(None).await; + *mgr.current_scopes.write().await = vec!["read".to_string()]; + + // round one: server challenges for "write" + let round_one = mgr.select_scopes(Some("write"), &[]); + assert_eq!(round_one, vec!["read".to_string(), "write".to_string()]); + *mgr.current_scopes.write().await = round_one; + + // round two: server challenges for "admin", earlier scopes are retained + let round_two = mgr.select_scopes(Some("admin"), &[]); + assert_eq!( + round_two, + vec!["read".to_string(), "write".to_string(), "admin".to_string()] + ); + } + + #[tokio::test] + async fn select_scopes_deduplicates_challenge_already_requested() { + let mgr = manager_with_metadata(None).await; + *mgr.current_scopes.write().await = vec!["read".to_string(), "write".to_string()]; + + let scopes = mgr.select_scopes(Some("write admin"), &[]); + + assert_eq!( + scopes, + vec!["read".to_string(), "write".to_string(), "admin".to_string()] + ); + } + + #[tokio::test] + async fn select_scopes_unions_resource_metadata_as_operational_requirement() { + let mgr = manager_with_metadata(None).await; + *mgr.current_scopes.write().await = vec!["read".to_string()]; + *mgr.resource_scopes.write().await = vec!["profile".to_string()]; + + let scopes = mgr.select_scopes(Some("write"), &[]); + + assert!(scopes.contains(&"read".to_string())); + assert!(scopes.contains(&"write".to_string())); + assert!(scopes.contains(&"profile".to_string())); + } + + #[test] + fn resolve_granted_scopes_uses_requested_when_response_omits_scope() { + let granted = AuthorizationManager::resolve_granted_scopes( + None, + &["read".to_string(), "write".to_string()], + &["read".to_string()], + ); + + assert_eq!(granted, vec!["read".to_string(), "write".to_string()]); + } + + #[test] + fn resolve_granted_scopes_honors_explicit_server_downgrade() { + let granted = AuthorizationManager::resolve_granted_scopes( + Some(vec!["read".to_string()]), + &["read".to_string(), "write".to_string()], + &["read".to_string()], + ); + + assert_eq!(granted, vec!["read".to_string()]); + } + + #[test] + fn resolve_granted_scopes_falls_back_to_current_when_nothing_requested() { + let granted = + AuthorizationManager::resolve_granted_scopes(None, &[], &["read".to_string()]); + + assert_eq!(granted, vec!["read".to_string()]); + } + #[tokio::test] async fn add_offline_access_if_supported_works_with_explicit_scopes() { let mgr = manager_with_metadata(Some(AuthorizationMetadata { @@ -4828,6 +8125,7 @@ mod tests { token_response: Some(make_token_response("my-access-token", Some(3600))), granted_scopes: vec![], token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -4845,6 +8143,7 @@ mod tests { token_response: Some(make_token_response("stale-token", Some(3600))), granted_scopes: vec![], token_received_at: Some(AuthorizationManager::now_epoch_secs() - 7200), + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -4863,6 +8162,7 @@ mod tests { token_response: Some(make_token_response("no-expiry-token", None)), granted_scopes: vec![], token_received_at: None, + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -4880,6 +8180,7 @@ mod tests { token_response: Some(make_token_response("almost-expired", Some(3600))), granted_scopes: vec![], token_received_at: Some(AuthorizationManager::now_epoch_secs() - 3590), + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -4898,6 +8199,7 @@ mod tests { token_response: Some(make_token_response("stale-token", Some(3600))), granted_scopes: vec![], token_received_at: Some(AuthorizationManager::now_epoch_secs() - 7200), + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -4989,6 +8291,86 @@ mod tests { // -- client credentials (SEP-1046) -- + #[cfg(feature = "auth-client-credentials-jwt")] + #[test] + fn client_authentication_audience_defaults_to_metadata_issuer() { + let metadata = AuthorizationMetadata { + authorization_endpoint: "https://auth.example.com/authorize".to_string(), + token_endpoint: "https://auth.example.com/token".to_string(), + issuer: Some("https://auth.example.com".to_string()), + ..Default::default() + }; + + let audience = super::client_authentication_audience(&metadata, None).unwrap(); + + assert_eq!(audience, "https://auth.example.com"); + } + + #[cfg(feature = "auth-client-credentials-jwt")] + #[test] + fn client_authentication_audience_honors_explicit_override() { + let metadata = AuthorizationMetadata { + authorization_endpoint: "https://auth.example.com/authorize".to_string(), + token_endpoint: "https://auth.example.com/token".to_string(), + issuer: Some("https://auth.example.com".to_string()), + ..Default::default() + }; + + let audience = super::client_authentication_audience( + &metadata, + Some("https://legacy.example.com/token"), + ) + .unwrap(); + + assert_eq!(audience, "https://legacy.example.com/token"); + } + + #[cfg(feature = "auth-client-credentials-jwt")] + #[test] + fn client_authentication_audience_rejects_missing_metadata_issuer() { + let metadata = AuthorizationMetadata { + authorization_endpoint: "https://auth.example.com/authorize".to_string(), + token_endpoint: "https://auth.example.com/token".to_string(), + ..Default::default() + }; + + let error = super::client_authentication_audience(&metadata, None).unwrap_err(); + + assert!(matches!(error, AuthError::ClientCredentialsError(_))); + } + + #[cfg(feature = "auth-client-credentials-jwt")] + #[test] + fn client_authentication_header_sets_explicit_type() { + let header = super::client_authentication_header(super::JwtSigningAlgorithm::ES256); + + assert_eq!(header.typ.as_deref(), Some("client-authentication+jwt")); + } + + #[cfg(feature = "auth-client-credentials-jwt")] + #[test] + fn client_credentials_endpoint_allows_http_between_loopback_hosts() { + let resource = Url::parse("http://localhost:8000/mcp").unwrap(); + let token_endpoint = Url::parse("http://127.0.0.1:9000/token").unwrap(); + + assert!(super::is_allowed_client_credentials_endpoint( + &resource, + &token_endpoint + )); + } + + #[cfg(feature = "auth-client-credentials-jwt")] + #[test] + fn client_credentials_endpoint_rejects_http_to_non_loopback_host() { + let resource = Url::parse("http://localhost:8000/mcp").unwrap(); + let token_endpoint = Url::parse("http://auth.example.com/token").unwrap(); + + assert!(!super::is_allowed_client_credentials_endpoint( + &resource, + &token_endpoint + )); + } + #[tokio::test] async fn configure_client_credentials_uses_request_body_auth_for_client_secret() { let mut mgr = manager_with_metadata(None).await; @@ -5003,6 +8385,40 @@ mod tests { assert!(matches!(oauth_client.auth_type(), AuthType::RequestBody)); } + #[test] + fn client_secret_credentials_debug_redacts_secret() { + let config = super::ClientCredentialsConfig::ClientSecret { + client_id: "client-id".to_string(), + client_secret: "client-secret-value".to_string(), + scopes: vec![], + resource: None, + }; + + let rendered = format!("{config:?}"); + + assert!(!rendered.contains("client-secret-value"), "{rendered}"); + } + + #[cfg(feature = "auth-client-credentials-jwt")] + #[test] + fn private_key_jwt_credentials_debug_redacts_signing_key() { + let config = super::ClientCredentialsConfig::PrivateKeyJwt { + client_id: "client-id".to_string(), + signing_key: b"private-signing-key-value".to_vec(), + signing_algorithm: super::JwtSigningAlgorithm::RS256, + token_endpoint_audience: None, + scopes: vec![], + resource: None, + }; + + let rendered = format!("{config:?}"); + + assert!( + !rendered.contains("private-signing-key-value"), + "{rendered}" + ); + } + #[tokio::test] async fn configure_client_credentials_sets_correct_client_id() { let mut mgr = manager_with_metadata(None).await; @@ -5159,6 +8575,53 @@ mod tests { resp } + async fn manager_with_refresh_error(error: &'static str) -> AuthorizationManager { + use axum::{Router, body::Body, http::Response, routing::post}; + + let app = Router::new().route( + "/token", + post(move || async move { + Response::builder() + .status(400) + .header("content-type", "application/json") + .body(Body::from( + serde_json::json!({ + "error": error, + "error_description": "refresh failed", + }) + .to_string(), + )) + .unwrap() + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + + let mut manager = manager_with_metadata(Some(AuthorizationMetadata { + authorization_endpoint: format!("http://{addr}/authorize"), + token_endpoint: format!("http://{addr}/token"), + ..Default::default() + })) + .await; + manager.configure_client(test_client_config()).unwrap(); + manager + .credential_store + .save(StoredCredentials { + client_id: "my-client".to_string(), + token_response: Some(make_token_response_with_refresh( + "old-token", + "my-refresh-token", + )), + granted_scopes: vec![], + token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: None, + }) + .await + .unwrap(); + manager + } + #[tokio::test] async fn refresh_token_returns_error_when_no_stored_credentials() { let mut manager = manager_with_metadata(None).await; @@ -5181,6 +8644,7 @@ mod tests { token_response: None, granted_scopes: vec![], token_received_at: None, + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -5201,13 +8665,38 @@ mod tests { token_response: Some(make_token_response("old-token", Some(3600))), granted_scopes: vec![], token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); let err = manager.refresh_token().await.unwrap_err(); + assert!( + matches!(err, AuthError::AuthorizationRequired), + "expected AuthorizationRequired when no refresh token, got: {err:?}" + ); + } + + #[tokio::test] + async fn invalid_grant_refresh_requires_reauthorization() { + let manager = manager_with_refresh_error("invalid_grant").await; + + let err = manager.try_refresh_or_reauth().await.unwrap_err(); + + assert!( + matches!(err, AuthError::AuthorizationRequired), + "expected AuthorizationRequired when the refresh token is rejected, got: {err:?}" + ); + } + + #[tokio::test] + async fn temporary_refresh_failure_does_not_require_reauthorization() { + let manager = manager_with_refresh_error("temporarily_unavailable").await; + + let err = manager.try_refresh_or_reauth().await.unwrap_err(); + assert!( matches!(err, AuthError::TokenRefreshFailed(_)), - "expected TokenRefreshFailed when no refresh token, got: {err:?}" + "expected TokenRefreshFailed for a temporary provider failure, got: {err:?}" ); } @@ -5267,6 +8756,7 @@ mod tests { )), granted_scopes: vec![], token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: None, }) .await .unwrap(); @@ -5472,6 +8962,7 @@ mod tests { )), granted_scopes: vec!["read".to_string(), "write".to_string()], token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -5489,6 +8980,43 @@ mod tests { assert_eq!(scope_parts, vec!["read", "write"]); } + #[tokio::test] + async fn refresh_token_includes_resource_parameter() { + let (base_url, captured) = start_token_server().await; + + let mut manager = manager_with_metadata(Some(AuthorizationMetadata { + authorization_endpoint: format!("{}/authorize", base_url), + token_endpoint: format!("{}/token", base_url), + ..Default::default() + })) + .await; + manager.configure_client(test_client_config()).unwrap(); + + let stored = StoredCredentials { + client_id: "my-client".to_string(), + token_response: Some(make_token_response_with_refresh( + "old-token", + "my-refresh-token", + )), + granted_scopes: vec![], + token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: None, + }; + manager.credential_store.save(stored).await.unwrap(); + + manager.refresh_token().await.unwrap(); + + let body = captured.lock().unwrap().take().unwrap(); + let params: std::collections::HashMap<_, _> = url::form_urlencoded::parse(body.as_bytes()) + .into_owned() + .collect(); + assert_eq!( + params.get("resource").map(String::as_str), + Some("http://localhost/"), + "refresh requests must carry the RFC 8707 resource parameter, got body: {body}" + ); + } + #[tokio::test] async fn refresh_token_adds_offline_access_when_as_supports_it() { let (base_url, captured) = start_token_server().await; @@ -5510,6 +9038,7 @@ mod tests { )), granted_scopes: vec!["read".to_string()], token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -5547,6 +9076,7 @@ mod tests { )), granted_scopes: vec![], token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -5585,6 +9115,7 @@ mod tests { )), granted_scopes: vec![], token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -5648,6 +9179,7 @@ mod tests { )), granted_scopes: vec![], token_received_at: Some(AuthorizationManager::now_epoch_secs()), + issuer: None, }; manager.credential_store.save(stored).await.unwrap(); @@ -5658,4 +9190,349 @@ mod tests { "a rotated refresh token from the response should replace the old one" ); } + + #[derive(Clone)] + struct RefreshStore { + credentials: InMemoryCredentialStore, + lock: Arc>, + events: Arc>>, + guard_requested: Arc, + save_started: Arc, + save_gate: Option>, + fail_at: Option<&'static str>, + } + + struct ObservedRefreshGuard { + _lock: OwnedMutexGuard<()>, + events: Arc>>, + } + + impl Drop for ObservedRefreshGuard { + fn drop(&mut self) { + self.events.lock().unwrap().push("release"); + } + } + + #[async_trait::async_trait] + impl CredentialStore for RefreshStore { + async fn load(&self) -> Result, AuthError> { + self.events.lock().unwrap().push("load"); + if self.fail_at == Some("load") { + return Err(AuthError::CredentialStoreError("load failed".into())); + } + self.credentials.load().await + } + + async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> { + self.events.lock().unwrap().push("save"); + self.save_started.add_permits(1); + if let Some(gate) = &self.save_gate { + gate.acquire().await.unwrap().forget(); + } + if self.fail_at == Some("save") { + return Err(AuthError::CredentialStoreError("save failed".into())); + } + self.credentials.save(credentials).await?; + self.events.lock().unwrap().push("saved"); + Ok(()) + } + + async fn clear(&self) -> Result<(), AuthError> { + self.credentials.clear().await + } + + async fn acquire_refresh_guard(&self) -> Result, AuthError> { + self.events.lock().unwrap().push("acquire"); + self.guard_requested.add_permits(1); + if self.fail_at == Some("guard") { + return Err(AuthError::CredentialStoreError("guard failed".into())); + } + let lock = self.lock.clone().lock_owned().await; + self.events.lock().unwrap().push("acquired"); + Ok(Some(CredentialRefreshGuard::new(ObservedRefreshGuard { + _lock: lock, + events: self.events.clone(), + }))) + } + } + + async fn refresh_store() -> RefreshStore { + let credentials = StoredCredentials::new( + "my-client".into(), + Some(make_token_response_with_refresh("old-token", "old-refresh")), + vec!["read".into()], + Some(AuthorizationManager::now_epoch_secs()), + ); + let credential_store = InMemoryCredentialStore::new(); + credential_store.save(credentials).await.unwrap(); + RefreshStore { + credentials: credential_store, + lock: Arc::new(Mutex::new(())), + events: Arc::new(StdMutex::new(Vec::new())), + guard_requested: Arc::new(Semaphore::new(0)), + save_started: Arc::new(Semaphore::new(0)), + save_gate: None, + fail_at: None, + } + } + + struct RefreshHttpClient { + recording: RecordingOAuthHttpClient, + events: Arc>>, + } + + impl OAuthHttpClient for RefreshHttpClient { + fn execute(&self, request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { + self.events.lock().unwrap().push("provider"); + self.recording.execute(request) + } + } + + fn refresh_http_client(store: &RefreshStore) -> Arc { + Arc::new(RefreshHttpClient { + recording: RecordingOAuthHttpClient::with_responses( + [ + ("new-token", "new-refresh"), + ("latest-token", "latest-refresh"), + ] + .into_iter() + .map(|(access, refresh)| { + http_response( + 200, + serde_json::json!({ + "access_token": access, "token_type": "Bearer", + "expires_in": 3600, "refresh_token": refresh + }), + ) + }) + .collect(), + ), + events: store.events.clone(), + }) + } + + async fn refresh_manager( + store: RefreshStore, + http_client: Arc, + ) -> AuthorizationManager { + let mut manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + http_client, + ) + .await + .unwrap(); + manager.set_metadata(AuthorizationMetadata { + authorization_endpoint: "https://auth.example.com/authorize".into(), + token_endpoint: "https://auth.example.com/token".into(), + ..Default::default() + }); + manager.configure_client(test_client_config()).unwrap(); + manager.set_credential_store(store); + *manager.current_scopes.write().await = vec!["cached".into()]; + manager + } + + async fn wait_for_permits(semaphore: &Semaphore, count: u32) { + tokio::time::timeout( + std::time::Duration::from_secs(5), + semaphore.acquire_many(count), + ) + .await + .unwrap() + .unwrap() + .forget(); + } + + #[tokio::test] + async fn refresh_guard_spans_load_exchange_and_completed_save() { + let store = refresh_store().await; + let manager = refresh_manager(store.clone(), refresh_http_client(&store)).await; + + manager.refresh_token().await.unwrap(); + + assert_eq!( + *store.events.lock().unwrap(), + [ + "acquire", "acquired", "load", "provider", "save", "saved", "release" + ] + ); + let saved = store.credentials.load().await.unwrap().unwrap(); + assert_eq!( + saved + .token_response + .unwrap() + .refresh_token() + .unwrap() + .secret(), + "new-refresh" + ); + assert_eq!(saved.granted_scopes, ["read"]); + assert!(store.lock.try_lock().is_ok()); + } + + #[tokio::test] + async fn concurrent_refreshes_wait_for_save_and_use_the_latest_token() { + let mut store = refresh_store().await; + let save_gate = Arc::new(Semaphore::new(0)); + store.save_gate = Some(save_gate.clone()); + let http_client = refresh_http_client(&store); + let first_manager = refresh_manager(store.clone(), http_client.clone()).await; + let second_manager = refresh_manager(store.clone(), http_client.clone()).await; + + let first = tokio::spawn(async move { first_manager.refresh_token().await }); + wait_for_permits(&store.save_started, 1).await; + let second = tokio::spawn(async move { second_manager.refresh_token().await }); + wait_for_permits(&store.guard_requested, 2).await; + assert_eq!(http_client.recording.requests().len(), 1); + assert!(store.lock.try_lock().is_err()); + + save_gate.add_permits(2); + let (first, second) = tokio::join!(first, second); + assert_eq!(first.unwrap().unwrap().access_token().secret(), "new-token"); + assert_eq!( + second.unwrap().unwrap().access_token().secret(), + "latest-token" + ); + let refresh_tokens: Vec = http_client + .recording + .requests() + .iter() + .map(|request| { + url::form_urlencoded::parse(&request.body) + .find_map(|(key, value)| (key == "refresh_token").then(|| value.into_owned())) + .unwrap() + }) + .collect(); + assert_eq!(refresh_tokens, ["old-refresh", "new-refresh"]); + let saved = store.credentials.load().await.unwrap().unwrap(); + assert_eq!( + saved + .token_response + .unwrap() + .refresh_token() + .unwrap() + .secret(), + "latest-refresh" + ); + assert!(store.lock.try_lock().is_ok()); + } + + #[tokio::test] + async fn refresh_rejects_credentials_for_another_client() { + let store = refresh_store().await; + let mut credentials = store.credentials.load().await.unwrap().unwrap(); + credentials.client_id = "other-client".into(); + store.credentials.save(credentials).await.unwrap(); + let http_client = refresh_http_client(&store); + let manager = refresh_manager(store.clone(), http_client.clone()).await; + + assert!(matches!( + manager.refresh_token().await, + Err(AuthError::AuthorizationRequired) + )); + assert!(http_client.recording.requests().is_empty()); + assert!(store.lock.try_lock().is_ok()); + } + + #[tokio::test] + async fn refresh_rejects_credentials_for_another_client_without_a_guard() { + let (base_url, captured) = start_token_server().await; + let mut manager = manager_with_metadata(Some(AuthorizationMetadata { + authorization_endpoint: format!("{base_url}/authorize"), + token_endpoint: format!("{base_url}/token"), + ..Default::default() + })) + .await; + manager.configure_client(test_client_config()).unwrap(); + manager + .credential_store + .save(StoredCredentials::new( + "other-client".into(), + Some(make_token_response_with_refresh("old-token", "old-refresh")), + vec!["read".into()], + Some(AuthorizationManager::now_epoch_secs()), + )) + .await + .unwrap(); + + let error = manager.refresh_token().await.unwrap_err(); + + assert!( + matches!(error, AuthError::AuthorizationRequired), + "a client mismatch must require reauthorization, got: {error:?}" + ); + assert!( + captured.lock().unwrap().is_none(), + "a client mismatch must be caught before the refresh token leaves the process" + ); + } + + #[tokio::test] + async fn refresh_without_a_guard_keeps_stored_scopes_when_response_omits_them() { + // start_token_server answers without a `scope`, matching a provider that + // grants the request in full. + let (base_url, _captured) = start_token_server().await; + let mut manager = manager_with_metadata(Some(AuthorizationMetadata { + authorization_endpoint: format!("{base_url}/authorize"), + token_endpoint: format!("{base_url}/token"), + ..Default::default() + })) + .await; + manager.configure_client(test_client_config()).unwrap(); + manager + .credential_store + .save(StoredCredentials::new( + "my-client".into(), + Some(make_token_response_with_refresh("old-token", "old-refresh")), + vec!["read".into()], + Some(AuthorizationManager::now_epoch_secs()), + )) + .await + .unwrap(); + *manager.current_scopes.write().await = vec!["stale".into()]; + + manager.refresh_token().await.unwrap(); + + let saved = manager.credential_store.load().await.unwrap().unwrap(); + assert_eq!( + saved.granted_scopes, + ["read"], + "the stored grant outranks the per-process scope cache" + ); + assert_eq!( + manager.get_current_scopes().await, + ["read"], + "the refreshed grant must replace the stale scope cache" + ); + } + + #[rstest] + #[case("guard", 0)] + #[case("load", 0)] + #[case("save", 1)] + #[tokio::test] + async fn refresh_store_failures_release_the_guard( + #[case] phase: &'static str, + #[case] provider_requests: usize, + ) { + let mut store = refresh_store().await; + store.fail_at = Some(phase); + let http_client = refresh_http_client(&store); + let manager = refresh_manager(store.clone(), http_client.clone()).await; + + assert!(matches!(manager.refresh_token().await, + Err(AuthError::CredentialStoreError(message)) if message == format!("{phase} failed"))); + assert_eq!(http_client.recording.requests().len(), provider_requests); + let saved = store.credentials.load().await.unwrap().unwrap(); + assert_eq!( + saved + .token_response + .unwrap() + .refresh_token() + .unwrap() + .secret(), + "old-refresh" + ); + assert!(store.lock.try_lock().is_ok()); + } } diff --git a/crates/rmcp/src/transport/auth/enterprise.rs b/crates/rmcp/src/transport/auth/enterprise.rs new file mode 100644 index 0000000..226c8d5 --- /dev/null +++ b/crates/rmcp/src/transport/auth/enterprise.rs @@ -0,0 +1,767 @@ +//! Non-interactive enterprise-managed authorization (EMA/XAA) token exchanges. +//! +//! Exchanges an enterprise refresh token for an ID-JAG, then for an MCP access +//! token. Callers must discover and approve both servers and their registrations +//! before supplying a credential. This module does not discover servers, log in, +//! persist credentials, or decide when to reauthenticate. +//! +//! Configure each pre-registered client's approved authentication method separately: +//! HTTP Basic, a client secret in the request body, or a freshly signed JWT client +//! assertion. Public-client authentication requires the server's explicit approval. +//! This helper requires one MCP resource and does not implement Rich +//! Authorization Requests or DPoP. Redemption consumes the SDK's ID-JAG handle, +//! without automatic retries; server-side replay policy remains the server's responsibility. +//! +//! ID-JAG checks below enforce structure and claim bindings, not cryptographic +//! signature verification. Assertions come directly from the trusted IdP token +//! endpoint; the resource authorization server must verify their signatures. +//! +//! ```no_run +//! use oauth2::{ClientSecret, RefreshToken}; +//! use rmcp::transport::auth::{default_oauth_http_client, enterprise::*}; +//! +//! # async fn authorize(refresh: &RefreshToken, idp_secret: ClientSecret, resource_secret: ClientSecret) -> Result<(), Box> { +//! // Enable `auth-enterprise-managed` and a TLS feature such as `reqwest`. +//! let http = default_oauth_http_client()?; +//! let token = EmaExchangeRequest::new( +//! EmaAuthorizationServer::new("https://idp.example", "https://idp.example/token", "idp-client") +//! .with_client_authentication(EmaClientAuthentication::ClientSecretBasic(idp_secret)), +//! EmaAuthorizationServer::new("https://as.example", "https://as.example/token", "mcp-client") +//! .with_client_authentication(EmaClientAuthentication::ClientSecretBasic(resource_secret)), +//! "https://mcp.example", refresh, +//! ).with_scopes(["files.read"]).exchange(&http, &http).await?; +//! // Use token.access_token only for the approved MCP resource; never log it. +//! // token.scopes contains the final granted scopes, which may be narrower. +//! # Ok(()) } +//! ``` + +use std::{ + collections::HashSet, + sync::Arc, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use base64::{ + Engine, + engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD}, +}; +use oauth2::{AccessToken, ClientSecret, RefreshToken}; +use serde::{Deserialize, de::DeserializeOwned}; +use thiserror::Error; +use url::{Host, Url}; + +use super::{ + DEFAULT_HTTP_TIMEOUT, MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES, OAuthHttpClient, + OAuthHttpRedirectPolicy, OAuthHttpRequest, +}; + +const ID_JAG_TOKEN_TYPE: &str = "urn:ietf:params:oauth:token-type:id-jag"; + +/// Authentication approved for a pre-registered client at one authorization server. +/// +/// The selected method is used as configured, without negotiation or fallback. +#[derive(Clone)] +#[non_exhaustive] +pub enum EmaClientAuthentication { + /// Public client (`token_endpoint_auth_method=none`), only if the server permits it. + None, + /// `client_secret_basic`, with OAuth form encoding before HTTP Basic encoding. + ClientSecretBasic(ClientSecret), + /// `client_secret_post`, for servers requiring credentials in the request body. + ClientSecretPost(ClientSecret), + /// Fresh JWT client assertions, such as `private_key_jwt` or `client_secret_jwt`. + /// Signing, key custody, claims, and the registered algorithm belong to the provider. + JwtAssertion(Arc), +} + +impl std::fmt::Debug for EmaClientAuthentication { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::None => "None", + Self::ClientSecretBasic(_) => "ClientSecretBasic { .. }", + Self::ClientSecretPost(_) => "ClientSecretPost { .. }", + Self::JwtAssertion(_) => "JwtAssertion { .. }", + }) + } +} + +impl EmaClientAuthentication { + fn validate(&self) -> Result<(), EmaError> { + if let Self::ClientSecretBasic(secret) | Self::ClientSecretPost(secret) = self + && secret.secret().trim().is_empty() + { + return Err(EmaError::InvalidRequest("client secret must not be empty")); + } + Ok(()) + } +} + +/// A signed JWT used to authenticate the client, distinct from the ID-JAG grant. +pub struct EmaClientAssertion(String); + +impl EmaClientAssertion { + /// Wrap a fresh assertion without exposing it through `Debug`. + pub fn new(assertion: String) -> Self { + Self(assertion) + } +} + +impl std::fmt::Debug for EmaClientAssertion { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("EmaClientAssertion { .. }") + } +} + +/// Creates client assertions on demand, allowing keys to remain in an external signer. +/// +/// Called once immediately before each token request, including delayed ID-JAG +/// redemption. Set `iss` and `sub` to the registered client identifier and `aud` +/// to the server's approved audience, with a short expiration and a fresh `jti`. +/// The SDK does not sign or validate these assertions. Provider failures are +/// sanitized and the call shares the token request's timeout. Cancellation may +/// occur when that deadline expires. +#[async_trait::async_trait] +pub trait EmaClientAssertionProvider: Send + Sync { + async fn create_assertion( + &self, + server: &EmaAuthorizationServer, + ) -> Result>; +} + +/// A trusted server with a pre-registered client and its approved authentication method. +#[derive(Clone)] +#[non_exhaustive] +pub struct EmaAuthorizationServer { + /// Exact issuer identifier from approved metadata. + pub issuer: String, + /// Token endpoint from that metadata. + pub token_endpoint: String, + /// Pre-registered client identifier. + pub client_id: String, + client_authentication: EmaClientAuthentication, +} + +impl EmaAuthorizationServer { + /// Use approved metadata and a public-client registration (`token_endpoint_auth_method=none`). + /// Set [`Self::with_client_authentication`] for a confidential client. + pub fn new( + issuer: impl Into, + token_endpoint: impl Into, + client_id: impl Into, + ) -> Self { + Self { + issuer: issuer.into(), + token_endpoint: token_endpoint.into(), + client_id: client_id.into(), + client_authentication: EmaClientAuthentication::None, + } + } + + /// Select the authentication method approved for this server's client registration. + pub fn with_client_authentication(mut self, authentication: EmaClientAuthentication) -> Self { + self.client_authentication = authentication; + self + } +} + +/// A refresh-token exchange bound to one MCP resource and two registered clients. +pub struct EmaExchangeRequest<'a> { + idp: EmaAuthorizationServer, + resource_as: EmaAuthorizationServer, + resource: &'a str, + refresh_token: &'a RefreshToken, + scopes: Vec, +} + +impl<'a> EmaExchangeRequest<'a> { + /// No scope parameter is sent until [`Self::with_scopes`] is used. + pub fn new( + idp: EmaAuthorizationServer, + resource_as: EmaAuthorizationServer, + resource: &'a str, + refresh_token: &'a RefreshToken, + ) -> Self { + Self { + idp, + resource_as, + resource, + refresh_token, + scopes: Vec::new(), + } + } + + /// Request distinct non-empty scope tokens; the IdP may narrow them. + /// An empty iterator omits `scope`, rather than requesting an empty grant. + pub fn with_scopes(mut self, scopes: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.scopes = scopes.into_iter().map(Into::into).collect(); + self + } + + /// Obtain an ID-JAG without redirects or retries, checking its resource/client bindings. + /// The returned assertion does not retain the enterprise refresh token. + pub async fn exchange_id_jag( + self, + idp_http: &dyn OAuthHttpClient, + ) -> Result { + for endpoint in [ + self.resource, + &self.idp.issuer, + &self.idp.token_endpoint, + &self.resource_as.issuer, + &self.resource_as.token_endpoint, + ] { + validate_endpoint(endpoint)?; + } + if self.idp.issuer == self.resource_as.issuer { + return Err(EmaError::InvalidRequest( + "IdP and resource AS issuers must differ", + )); + } + if self.idp.client_id.trim().is_empty() + || self.resource_as.client_id.trim().is_empty() + || self.refresh_token.secret().trim().is_empty() + { + return Err(EmaError::InvalidRequest( + "client IDs and refresh token must not be empty", + )); + } + self.idp.client_authentication.validate()?; + self.resource_as.client_authentication.validate()?; + let requested: HashSet<&str> = self.scopes.iter().map(String::as_str).collect(); + if requested.len() != self.scopes.len() || self.scopes.iter().any(|s| !is_scope_token(s)) { + return Err(EmaError::InvalidRequest( + "scopes must be distinct non-empty tokens", + )); + } + let mut params = vec![ + ( + "grant_type", + "urn:ietf:params:oauth:grant-type:token-exchange", + ), + ("requested_token_type", ID_JAG_TOKEN_TYPE), + ( + "subject_token_type", + "urn:ietf:params:oauth:token-type:refresh_token", + ), + ("subject_token", self.refresh_token.secret()), + ("audience", self.resource_as.issuer.as_str()), + ("resource", self.resource), + ]; + let scope = self.scopes.join(" "); + if !scope.is_empty() { + params.push(("scope", &scope)); + } + let jag: IdJagResponse = post_form( + idp_http, + &self.idp, + ¶ms, + EmaExchangeStage::IdentityProvider, + None, + unix_time, + ) + .await?; + let (granted, expires_at) = jag.validate(&self, &requested)?; + Ok(EmaIdJag { + assertion: AccessToken::new(jag.access_token), + scopes: granted, + resource_as: self.resource_as, + resource: self.resource.to_owned(), + expires_at, + }) + } + + /// Perform both exchanges with independently routed HTTP clients and no automatic retries. + pub async fn exchange( + self, + idp_http: &dyn OAuthHttpClient, + resource_http: &dyn OAuthHttpClient, + ) -> Result { + self.exchange_id_jag(idp_http) + .await? + .exchange(resource_http) + .await + } +} + +/// An IdP-issued ID-JAG whose structure and bindings have been checked, not its signature. +pub struct EmaIdJag { + assertion: AccessToken, + scopes: HashSet, + resource_as: EmaAuthorizationServer, + resource: String, + expires_at: u64, +} + +impl std::fmt::Debug for EmaAuthorizationServer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("EmaAuthorizationServer { .. }") + } +} + +impl std::fmt::Debug for EmaExchangeRequest<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("EmaExchangeRequest { .. }") + } +} + +impl std::fmt::Debug for EmaIdJag { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("EmaIdJag { .. }") + } +} + +impl EmaIdJag { + /// The assertion to present to the approved resource authorization server. + pub fn assertion(&self) -> &AccessToken { + &self.assertion + } + + /// The scope tokens carried by the assertion; an empty set means scope was omitted. + pub fn scopes(&self) -> &HashSet { + &self.scopes + } + + /// Redeem this assertion once, at its approved resource AS, without redirects or retries. + /// Consume the grant so this helper cannot accidentally replay it after a failed exchange. + pub async fn exchange(self, http: &dyn OAuthHttpClient) -> Result { + self.exchange_with_clock(http, unix_time).await + } + + async fn exchange_with_clock( + self, + http: &dyn OAuthHttpClient, + now: impl Fn() -> Result + Sync, + ) -> Result { + if self.expires_at <= now()? { + return Err(EmaError::InvalidRequest("ID-JAG expired before redemption")); + } + // Only the assertion carries authority: repeating resource/scope could undo narrowing. + let token: ResourceTokenResponse = post_form( + http, + &self.resource_as, + &[ + ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"), + ("assertion", self.assertion.secret()), + ], + EmaExchangeStage::ResourceAuthorizationServer, + Some(self.expires_at), + now, + ) + .await?; + token.validate(&self.resource, self.scopes) + } +} + +/// A resource-bound bearer with secret-safe diagnostics and its optional lifetime. +#[derive(Clone)] +#[non_exhaustive] +pub struct EmaAccessToken { + pub access_token: AccessToken, + pub expires_in: Option, + /// Resource-AS scopes, or the ID-JAG scopes when the response omits `scope`. + /// Empty when both the ID-JAG and response omit scopes. + pub scopes: HashSet, +} + +impl std::fmt::Debug for EmaAccessToken { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("EmaAccessToken { .. }") + } +} + +/// The endpoint that failed, allowing the caller to apply its own credential lifecycle policy. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum EmaExchangeStage { + IdentityProvider, + ResourceAuthorizationServer, +} + +/// Sanitized failures. Raw HTTP adapter errors and provider response bodies are never retained. +#[derive(Debug, Error, PartialEq, Eq)] +#[non_exhaustive] +pub enum EmaError { + #[error("invalid EMA exchange request: {0}")] + InvalidRequest(&'static str), + #[error("invalid EMA response from {stage:?}: {message}")] + InvalidResponse { + stage: EmaExchangeStage, + message: &'static str, + }, + #[error("EMA request to {0:?} failed")] + RequestFailed(EmaExchangeStage), + #[error("{0:?}: invalid_grant")] + InvalidGrant(EmaExchangeStage), + #[error("{0:?}: insufficient_user_authentication")] + InsufficientUserAuthentication(EmaExchangeStage), + #[error("{stage:?} returned HTTP {status}: {code}")] + OAuthRejected { + stage: EmaExchangeStage, + status: u16, + code: &'static str, + }, +} + +async fn post_form( + http: &dyn OAuthHttpClient, + server: &EmaAuthorizationServer, + params: &[(&str, &str)], + stage: EmaExchangeStage, + grant_expires_at: Option, + now: impl Fn() -> Result + Sync, +) -> Result { + let response = tokio::time::timeout(DEFAULT_HTTP_TIMEOUT, async { + // Generate assertions at the request boundary, not when server configuration is built. + let client_assertion = match &server.client_authentication { + EmaClientAuthentication::JwtAssertion(provider) => { + let assertion = provider + .create_assertion(server) + .await + .map_err(|_| EmaError::RequestFailed(stage))?; + if assertion.0.trim().is_empty() { + return Err(EmaError::InvalidRequest( + "client assertion must not be empty", + )); + } + Some(assertion) + } + _ => None, + }; + let request = { + let mut form = url::form_urlencoded::Serializer::new(String::new()); + form.extend_pairs(params.iter().copied()); + let mut request = oauth2::http::Request::builder() + .method("POST") + .uri(&server.token_endpoint) + .header("content-type", "application/x-www-form-urlencoded") + .header("accept", "application/json"); + match &server.client_authentication { + EmaClientAuthentication::ClientSecretBasic(secret) => { + let client_id: String = + url::form_urlencoded::byte_serialize(server.client_id.as_bytes()).collect(); + let secret: String = + url::form_urlencoded::byte_serialize(secret.secret().as_bytes()).collect(); + let encoded = STANDARD.encode(format!("{client_id}:{secret}")); + let mut header = + oauth2::http::HeaderValue::from_str(&format!("Basic {encoded}")).map_err( + |_| EmaError::InvalidRequest("invalid client authentication header"), + )?; + header.set_sensitive(true); + request = request.header(oauth2::http::header::AUTHORIZATION, header); + } + EmaClientAuthentication::ClientSecretPost(secret) => { + form.append_pair("client_id", &server.client_id) + .append_pair("client_secret", secret.secret()); + } + EmaClientAuthentication::None | EmaClientAuthentication::JwtAssertion(_) => { + form.append_pair("client_id", &server.client_id); + } + } + if let Some(assertion) = client_assertion { + form.append_pair( + "client_assertion_type", + "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + ) + .append_pair("client_assertion", &assertion.0); + } + request + .body(form.finish().into_bytes()) + .map_err(|_| EmaError::InvalidRequest("invalid token endpoint URI"))? + }; + // An external signer may outlive the grant even when it meets the request deadline. + if let Some(expires_at) = grant_expires_at + && expires_at <= now()? + { + return Err(EmaError::InvalidRequest("ID-JAG expired before redemption")); + } + http.execute(OAuthHttpRequest::new( + request, + OAuthHttpRedirectPolicy::Stop, + )) + .await + .map_err(|_| EmaError::RequestFailed(stage)) + }) + .await + .map_err(|_| EmaError::RequestFailed(stage))??; + let invalid = |message| EmaError::InvalidResponse { stage, message }; + if response.body().len() > MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES { + return Err(invalid("response body too large")); + } + if !response.status().is_success() { + #[derive(Deserialize)] + struct OAuthError { + error: Option, + } + let error = serde_json::from_slice::(response.body()).ok(); + let code = match error.as_ref().and_then(|e| e.error.as_deref()) { + Some("invalid_grant") => return Err(EmaError::InvalidGrant(stage)), + Some("insufficient_user_authentication") => { + return Err(EmaError::InsufficientUserAuthentication(stage)); + } + Some("invalid_request") => "invalid_request", + Some("invalid_client") => "invalid_client", + Some("invalid_scope") => "invalid_scope", + Some("invalid_target") => "invalid_target", + Some("unauthorized_client") => "unauthorized_client", + Some("unsupported_grant_type") => "unsupported_grant_type", + Some("access_denied") => "access_denied", + Some("temporarily_unavailable") => "temporarily_unavailable", + Some("server_error") => "server_error", + _ => "OAuth token request rejected", + }; + return Err(EmaError::OAuthRejected { + stage, + status: response.status().as_u16(), + code, + }); + } + serde_json::from_slice(response.body()).map_err(|_| invalid("malformed token response")) +} + +fn validate_endpoint(value: &str) -> Result<(), EmaError> { + let url = Url::parse(value).map_err(|_| EmaError::InvalidRequest("invalid endpoint URL"))?; + let loopback = match url.host() { + Some(Host::Domain(host)) => host.eq_ignore_ascii_case("localhost"), + Some(Host::Ipv4(ip)) => ip.is_loopback(), + Some(Host::Ipv6(ip)) => ip.is_loopback(), + None => false, + }; + if (url.scheme() != "https" && !(url.scheme() == "http" && loopback)) + || !url.username().is_empty() + || url.password().is_some() + || url.fragment().is_some() + { + return Err(EmaError::InvalidRequest( + "endpoint must use HTTPS or HTTP loopback without userinfo or fragments", + )); + } + Ok(()) +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum Resource { + Single(String), + Multiple(Vec), +} + +impl Resource { + fn is_exact(&self, expected: &str) -> bool { + match self { + Self::Single(value) => value == expected, + Self::Multiple(values) => values.as_slice() == [expected], + } + } +} + +#[derive(Deserialize)] +struct JwtHeader { + alg: String, + typ: Option, +} + +#[derive(Deserialize)] +struct IdJagClaims { + iss: String, + sub: String, + aud: Resource, + client_id: String, + jti: String, + exp: u64, + iat: u64, + resource: Resource, + scope: Option, + authorization_details: Option>, +} + +#[derive(Deserialize)] +struct IdJagResponse { + access_token: String, + issued_token_type: String, + token_type: String, + resource: Option, + scope: Option, + refresh_token: Option, + authorization_details: Option>, +} + +impl IdJagResponse { + fn validate( + &self, + request: &EmaExchangeRequest<'_>, + requested: &HashSet<&str>, + ) -> Result<(HashSet, u64), EmaError> { + let invalid = |message| EmaError::InvalidResponse { + stage: EmaExchangeStage::IdentityProvider, + message, + }; + if self.issued_token_type != ID_JAG_TOKEN_TYPE + || self.token_type != "N_A" + || self.refresh_token.is_some() + { + return Err(invalid("unsupported ID-JAG token type or refresh token")); + } + let mut parts = self.access_token.split('.'); + let (Some(header), Some(payload), Some(signature), None) = + (parts.next(), parts.next(), parts.next(), parts.next()) + else { + return Err(invalid("ID-JAG must be a compact signed JWT")); + }; + if header.is_empty() || payload.is_empty() || signature.is_empty() { + return Err(invalid("ID-JAG contains an empty JWT segment")); + } + let decode = |value| { + URL_SAFE_NO_PAD + .decode(value) + .map_err(|_| invalid("malformed ID-JAG encoding")) + }; + decode(signature)?; + let header: JwtHeader = serde_json::from_slice(&decode(header)?) + .map_err(|_| invalid("malformed ID-JAG header"))?; + let claims: IdJagClaims = serde_json::from_slice(&decode(payload)?) + .map_err(|_| invalid("malformed ID-JAG claims"))?; + if [&self.authorization_details, &claims.authorization_details] + .into_iter() + .any(|details| details.as_ref().is_some_and(|details| !details.is_empty())) + { + return Err(invalid("authorization_details is not supported")); + } + if header.alg.trim().is_empty() + || header.alg.eq_ignore_ascii_case("none") + || header.typ.as_deref() != Some("oauth-id-jag+jwt") + || claims.iss != request.idp.issuer + || !claims.aud.is_exact(&request.resource_as.issuer) + || claims.client_id != request.resource_as.client_id + || claims.sub.trim().is_empty() + || claims.jti.trim().is_empty() + { + return Err(invalid( + "ID-JAG type, issuer, audience, client, subject, or JWT ID mismatch", + )); + } + let now = unix_time()?; + if claims.exp <= now || claims.iat > now.saturating_add(60) { + return Err(invalid("expired or future-issued ID-JAG")); + } + if !claims.resource.is_exact(request.resource) + || self + .resource + .as_ref() + .is_some_and(|r| !r.is_exact(request.resource)) + { + return Err(invalid("ID-JAG resource mismatch")); + } + let parse = + |scope| parse_scope(scope).ok_or_else(|| invalid("malformed or duplicate scopes")); + let granted = match claims.scope.as_deref() { + Some(scope) => parse(scope)?, + None if requested.is_empty() => HashSet::new(), + None => return Err(invalid("ID-JAG is missing requested scope authorization")), + }; + if !requested.is_empty() && !granted.is_subset(requested) { + return Err(invalid("ID-JAG scope exceeds the request")); + } + match self.scope.as_deref() { + Some(scope) if parse(scope)? != granted => { + return Err(invalid("response scope differs from ID-JAG scope")); + } + None if !requested.is_empty() && granted != *requested => { + return Err(invalid("response omitted narrowed scope")); + } + _ => {} + } + Ok((granted.into_iter().map(str::to_owned).collect(), claims.exp)) + } +} + +fn parse_scope(scope: &str) -> Option> { + let scopes: HashSet<_> = scope.split(' ').collect(); + (scopes.iter().all(|s| is_scope_token(s)) && scopes.len() == scope.split(' ').count()) + .then_some(scopes) +} + +fn is_scope_token(scope: &str) -> bool { + !scope.is_empty() + && scope + .bytes() + .all(|b| matches!(b, b'!' | b'#'..=b'[' | b']'..=b'~')) +} + +fn unix_time() -> Result { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs()) + .map_err(|_| EmaError::InvalidRequest("system clock precedes the Unix epoch")) +} + +#[derive(Deserialize)] +struct ResourceTokenResponse { + access_token: String, + token_type: String, + expires_in: Option, + resource: Option, + scope: Option, + refresh_token: Option, + authorization_details: Option>, +} + +impl ResourceTokenResponse { + fn validate( + self, + resource: &str, + granted: HashSet, + ) -> Result { + let invalid = |message| EmaError::InvalidResponse { + stage: EmaExchangeStage::ResourceAuthorizationServer, + message, + }; + if self + .authorization_details + .as_ref() + .is_some_and(|details| !details.is_empty()) + { + return Err(invalid("authorization_details is not supported")); + } + if !self.token_type.eq_ignore_ascii_case("bearer") + || self.access_token.trim().is_empty() + || self.refresh_token.is_some() + || self.expires_in == Some(0) + { + return Err(invalid( + "invalid bearer token, lifetime, or unexpected refresh token", + )); + } + // The resource need not be echoed, but must agree with the ID-JAG if present. + if self + .resource + .as_ref() + .is_some_and(|r| !r.is_exact(resource)) + { + return Err(invalid("access token resource mismatch")); + } + // An omitted scope retains the authority carried by the assertion. + let scopes = if let Some(scope) = self.scope.as_deref() { + let scopes = + parse_scope(scope).ok_or_else(|| invalid("malformed or duplicate scopes"))?; + if !scopes.iter().all(|s| granted.contains(*s)) { + return Err(invalid("access token scope exceeds ID-JAG scope")); + } + scopes.into_iter().map(str::to_owned).collect() + } else { + granted + }; + Ok(EmaAccessToken { + access_token: AccessToken::new(self.access_token), + expires_in: self.expires_in.map(Duration::from_secs), + scopes, + }) + } +} + +#[cfg(test)] +#[path = "enterprise_tests.rs"] +mod tests; diff --git a/crates/rmcp/src/transport/auth/enterprise_tests.rs b/crates/rmcp/src/transport/auth/enterprise_tests.rs new file mode 100644 index 0000000..31a63c1 --- /dev/null +++ b/crates/rmcp/src/transport/auth/enterprise_tests.rs @@ -0,0 +1,1145 @@ +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, + time::{Duration, SystemTime}, +}; + +use base64::{ + Engine, + engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD}, +}; +use oauth2::{ClientSecret, HttpResponse, RefreshToken}; +use serde_json::{Value, json}; + +use super::{ + EmaExchangeStage::{IdentityProvider as Idp, ResourceAuthorizationServer as ResourceServer}, + *, +}; +use crate::transport::auth::{OAuthHttpClientError, OAuthHttpClientFuture, OAuthHttpRequest}; + +const IDP: &str = "https://idp.example?private-query"; +const AS: &str = "https://as.example?private-query"; +const RESOURCE: &str = "https://mcp.example?private-query"; +const IDP_TOKEN: &str = "https://idp.example/token?private-query"; +const AS_TOKEN: &str = "https://as.example/token?private-query"; +const BAD_SCOPES: &[&str] = &["", "files\tread", "\"", "\\", "\0", "读"]; + +#[derive(Default)] +struct MockHttp { + requests: Mutex>, + response: Mutex>>, +} + +impl MockHttp { + fn new(status: u16, body: Value) -> Self { + Self { + response: Mutex::new(Some(Ok(oauth2::http::Response::builder() + .status(status) + .body(serde_json::to_vec(&body).unwrap()) + .unwrap()))), + ..Self::default() + } + } +} + +impl OAuthHttpClient for MockHttp { + fn execute(&self, request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { + self.requests.lock().unwrap().push(request); + let response = self.response.lock().unwrap().take(); + let response = response.expect("unexpected HTTP"); + Box::pin(async move { response }) + } +} + +fn jwt(header: Value, claims: &Value) -> String { + format!( + "{}.{}.{}", + URL_SAFE_NO_PAD.encode(header.to_string()), + URL_SAFE_NO_PAD.encode(claims.to_string()), + URL_SAFE_NO_PAD.encode(b"synthetic-signature") + ) +} + +fn claims() -> Value { + let now = SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + json!({"iss":IDP,"aud":AS,"sub":"user","client_id":"mcp", + "jti":"jag-id","iat":now,"exp":now + 3600,"resource":RESOURCE,"scope":"files.read"}) +} + +fn jag(claims: &Value) -> Value { + let mut body = json!({"access_token":jwt(json!({"alg":"ES256","typ":"oauth-id-jag+jwt"}), claims), + "issued_token_type":ID_JAG_TOKEN_TYPE,"token_type":"N_A","resource":claims["resource"]}); + if let Some(scope) = claims.get("scope") { + body["scope"] = scope.clone(); + } + body +} + +async fn exchange(idp: &MockHttp, scopes: &str) -> Result { + let scopes = scopes + .split_ascii_whitespace() + .map(str::to_owned) + .collect::>(); + let refresh = RefreshToken::new("refresh-token".into()); + let resource_as = EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp"); + assert!(!format!("{resource_as:?}").contains("private-query")); + let request = EmaExchangeRequest::new( + EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp"), + resource_as, + RESOURCE, + &refresh, + ) + .with_scopes(scopes); + assert!(!format!("{request:?}").contains(refresh.secret())); + assert!(!format!("{request:?}").contains("private-query")); + let future = request.exchange_id_jag(idp); + fn is_send(_: &T) {} + is_send(&future); + future.await +} + +fn form(request: &OAuthHttpRequest) -> BTreeMap { + assert!(!request.request.headers().contains_key("authorization")); + authenticated_form(request) +} + +fn authenticated_form(request: &OAuthHttpRequest) -> BTreeMap { + assert_eq!(request.request.method(), "POST"); + let headers = request.request.headers(); + assert_eq!(headers["content-type"], "application/x-www-form-urlencoded"); + assert_eq!(headers["accept"], "application/json"); + assert_eq!(request.redirect_policy, OAuthHttpRedirectPolicy::Stop); + assert_eq!(request.timeout, Some(Duration::from_secs(30))); + let pairs = url::form_urlencoded::parse(request.request.body()) + .into_owned() + .collect::>(); + let fields = pairs.iter().cloned().collect::>(); + assert_eq!(pairs.len(), fields.len(), "duplicate form fields"); + fields +} + +#[tokio::test] +async fn refresh_exchange_preserves_exact_forms_and_signed_narrowing() { + for scopes in ["files.read files.write", ""] { + let response = jag(&claims()); + let idp = MockHttp::new(200, response.clone()); + let result = exchange(&idp, scopes).await.unwrap(); + assert_eq!( + result.assertion().secret(), + response["access_token"].as_str().unwrap() + ); + assert_eq!(result.scopes(), &HashSet::from(["files.read".to_owned()])); + assert!(!format!("{result:?}").contains(result.assertion().secret())); + assert!(!format!("{result:?}").contains("private-query")); + let requests = idp.requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].request.uri(), IDP_TOKEN); + let mut fields = form(&requests[0]); + assert_eq!( + fields.remove("scope"), + (!scopes.is_empty()).then(|| scopes.to_owned()) + ); + assert_eq!( + serde_json::to_value(fields).unwrap(), + json!({ + "grant_type":"urn:ietf:params:oauth:grant-type:token-exchange", + "requested_token_type":ID_JAG_TOKEN_TYPE,"subject_token":"refresh-token", + "subject_token_type":"urn:ietf:params:oauth:token-type:refresh_token", + "audience":AS,"resource":RESOURCE,"client_id":"idp" + }) + ); + } +} + +#[tokio::test] +async fn scopes_may_be_omitted_but_never_widened() { + for (requested, signed, echoed, valid) in [ + ("", Some("files.read"), None, true), + ("", None, None, true), + ("files.read", Some("files.read"), None, true), + ("files.read files.write", Some("files.read"), None, false), + ( + "files.read", + Some("files.read files.write"), + Some("files.read files.write"), + false, + ), + ("files.read", None, None, false), + ("", Some("files.read"), Some("files.write"), false), + ("", Some(" \t"), None, false), + ("", Some("files.read files.read"), None, false), + ("", Some("files.read"), Some(" \t"), false), + ("", Some("files.read"), Some("files.read files.read"), false), + ] { + let mut claims = claims(); + claims.as_object_mut().unwrap().remove("scope"); + if let Some(scope) = signed { + claims["scope"] = json!(scope); + } + let mut response = jag(&claims); + response.as_object_mut().unwrap().remove("scope"); + if let Some(scope) = echoed { + response["scope"] = json!(scope); + } + let result = exchange(&MockHttp::new(200, response), requested).await; + assert_eq!( + result.is_ok(), + valid, + "requested={requested:?}, signed={signed:?}, echoed={echoed:?}" + ); + if let Ok(result) = result { + let expected = signed.into_iter().map(str::to_owned).collect(); + assert_eq!(result.scopes(), &expected); + } + } +} + +#[tokio::test] +async fn unsupported_authorization_details_are_rejected_at_each_stage() { + const SECRET: &str = "authorization-details-secret"; + for (location, stage) in [ + ("claims", Idp), + ("idp response", Idp), + ("resource response", ResourceServer), + ] { + for (case, details, valid) in [ + ("absent", None, true), + ("null", Some(Value::Null), true), + ("empty", Some(json!([])), true), + ( + "additional authority", + Some( + json!([{"type":SECRET,"locations":["https://other.example"],"actions":["write"]}]), + ), + false, + ), + ("invalid member", Some(json!([null])), false), + ("object", Some(json!({"type":SECRET})), false), + ("string", Some(json!(SECRET)), false), + ] { + let mut claims = claims(); + let mut response = bearer(); + // Unrelated extension fields remain compatible at every boundary. + claims["vendor_extension"] = json!(SECRET); + response["vendor_extension"] = json!(SECRET); + if location == "claims" + && let Some(details) = &details + { + claims["authorization_details"] = details.clone(); + } + let mut idp_response = jag(&claims); + idp_response["vendor_extension"] = json!(SECRET); + if let Some(details) = details { + match location { + "idp response" => idp_response["authorization_details"] = details, + "resource response" => response["authorization_details"] = details, + _ => {} + } + } + let idp = MockHttp::new(200, idp_response); + let resource = MockHttp::new(200, response); + let result = EmaExchangeRequest::new( + EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp"), + EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp"), + RESOURCE, + &RefreshToken::new("refresh-token".into()), + ) + .with_scopes(["files.read"]) + .exchange(&idp, &resource) + .await; + assert_eq!(result.is_ok(), valid, "{location}: {case}"); + if let Err(error) = result { + assert!( + matches!(error, EmaError::InvalidResponse { stage: actual, .. } if actual == stage) + ); + assert!(!format!("{error:?} {error}").contains(SECRET)); + } + assert_eq!(idp.requests.lock().unwrap().len(), 1); + assert_eq!( + resource.requests.lock().unwrap().len(), + usize::from(valid || stage == ResourceServer), + "{location}: {case}" + ); + } + } +} + +#[tokio::test] +async fn invalid_jags_never_escape_validation() { + let original = claims(); + let mut cases = Vec::new(); + let invalid = json!({"iss":"https://other.example","aud":[AS,"other"], + "client_id":"other","sub":"","jti":" \t","exp":0,"iat":u64::MAX,"resource":[RESOURCE,"other"]}); + for (key, value) in invalid.as_object().unwrap() { + let mut changed = original.clone(); + changed[key] = value.clone(); + cases.push(jag(&changed)); + } + for scope in BAD_SCOPES { + let mut changed = original.clone(); + changed["scope"] = json!(scope); + cases.push(jag(&changed)); + } + for header in [ + json!({"alg":"ES256","typ":"JWT"}), + json!({"alg":"ES256"}), + json!({"alg":"none","typ":"oauth-id-jag+jwt"}), + ] { + let mut response = jag(&original); + response["access_token"] = json!(jwt(header, &original)); + cases.push(response); + } + let invalid = json!({"issued_token_type":"Bearer","token_type":"Bearer", + "refresh_token":"unsupported","resource":"https://other.example","access_token":"a.b.c.d"}); + for (key, value) in invalid.as_object().unwrap() { + let mut response = jag(&original); + response[key] = value.clone(); + cases.push(response); + } + for signature in ["", "signature", "not+base64url", "c2ln="] { + let mut response = jag(&original); + let assertion = response["access_token"].as_str().unwrap(); + let (signed, _) = assertion.rsplit_once('.').unwrap(); + response["access_token"] = json!(format!("{signed}.{signature}")); + cases.push(response); + } + for response in cases { + let result = exchange(&MockHttp::new(200, response), "files.read").await; + assert!(matches!( + result, + Err(EmaError::InvalidResponse { stage: Idp, .. }) + )); + } +} + +#[tokio::test] +async fn invalid_inputs_fail_before_http() { + // Each server occupies issuer, token endpoint, and client ID slots. + let original = [IDP, IDP_TOKEN, "idp", AS, AS_TOKEN, "mcp", RESOURCE, "rt"]; + let mut cases = Vec::new(); + for index in [0, 1, 3, 4, 6] { + for value in [ + "invalid", + "http://idp.example/token", + "https://user:pass@idp.example/token", + "https://idp.example/token#fragment", + ] { + let mut fields = original; + fields[index] = value; + cases.push((fields, vec![])); + } + } + for (index, value) in [(0, AS), (2, " "), (5, ""), (7, " \t")] { + let mut fields = original; + fields[index] = value; + cases.push((fields, vec![])); + } + for scope in BAD_SCOPES.iter().copied().chain(["files.read files.write"]) { + cases.push((original, vec![scope])); + } + cases.push((original, vec!["files.read", "files.read"])); + for (fields, scopes) in cases { + let http = MockHttp::default(); + let result = EmaExchangeRequest::new( + EmaAuthorizationServer::new(fields[0], fields[1], fields[2]), + EmaAuthorizationServer::new(fields[3], fields[4], fields[5]), + fields[6], + &RefreshToken::new(fields[7].into()), + ) + .with_scopes(scopes) + .exchange_id_jag(&http) + .await; + assert!(matches!(result, Err(EmaError::InvalidRequest(_)))); + assert!(http.requests.lock().unwrap().is_empty()); + } +} + +#[tokio::test] +async fn errors_and_redirects_cannot_reflect_credentials() { + const SECRET: &str = "secret-error-sentinel"; + for (status, code) in [ + (400, "invalid_grant"), + (400, "insufficient_user_authentication"), + (400, "invalid_client"), + (400, SECRET), + (302, SECRET), + ] { + let failure = json!({"error":code,"error_description":SECRET}); + let error = exchange(&MockHttp::new(status, failure), "") + .await + .unwrap_err(); + match code { + "invalid_grant" => assert_eq!(error, EmaError::InvalidGrant(Idp)), + "insufficient_user_authentication" => { + assert_eq!(error, EmaError::InsufficientUserAuthentication(Idp)) + } + _ => assert!( + matches!(error, EmaError::OAuthRejected {stage: Idp, status: actual, ..} if actual == status) + ), + } + assert!(!format!("{error:?} {error}").contains(SECRET)); + } + for adapter_failure in [false, true] { + let failure = MockHttp::new( + 200, + json!({"access_token":SECRET,"issued_token_type":SECRET}), + ); + if adapter_failure { + *failure.response.lock().unwrap() = Some(Err(SECRET.into())); + } + let error = exchange(&failure, "").await.unwrap_err(); + assert!(!format!("{error:?} {error}").contains(SECRET)); + assert!(std::error::Error::source(&error).is_none()); + if adapter_failure { + assert_eq!(error, EmaError::RequestFailed(Idp)); + } + } + let mut oversized = jag(&claims()); + oversized["ignored"] = json!("x".repeat(1024 * 1024)); + let error = exchange(&MockHttp::new(200, oversized), "") + .await + .unwrap_err(); + assert!(matches!( + error, + EmaError::InvalidResponse { stage: Idp, .. } + )); +} + +fn bearer() -> Value { + json!({"access_token":"resource-token","token_type":"Bearer","expires_in":300}) +} + +async fn redeem(http: &MockHttp) -> Result { + exchange( + &MockHttp::new(200, jag(&claims())), + "files.read files.write", + ) + .await? + .exchange(http) + .await +} + +#[tokio::test] +async fn full_exchange_uses_separate_clients_and_only_the_narrowed_assertion() { + for valid in [true, false] { + let mut claims = claims(); + if !valid { + claims["client_id"] = json!("other"); + } + let response = jag(&claims); + let idp = MockHttp::new(200, response.clone()); + let resource_as = MockHttp::new(200, bearer()); + let refresh = RefreshToken::new("refresh-token".into()); + let request = EmaExchangeRequest::new( + EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp"), + EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp"), + RESOURCE, + &refresh, + ) + .with_scopes(["files.read", "files.write"]); + let future = request.exchange(&idp, &resource_as); + fn is_send(_: &T) {} + is_send(&future); + let result = future.await; + assert_eq!(idp.requests.lock().unwrap().len(), 1); + let requests = resource_as.requests.lock().unwrap(); + assert_eq!(requests.len(), usize::from(valid)); + if !valid { + assert!(matches!( + result, + Err(EmaError::InvalidResponse { stage: Idp, .. }) + )); + continue; + } + let token = result.unwrap(); + assert_eq!(token.access_token.secret(), "resource-token"); + assert_eq!(token.expires_in, Some(Duration::from_secs(300))); + assert_eq!(token.scopes, HashSet::from(["files.read".to_owned()])); + assert!(!format!("{token:?}").contains(token.access_token.secret())); + assert_eq!(requests[0].request.uri(), AS_TOKEN); + assert_eq!( + serde_json::to_value(form(&requests[0])).unwrap(), + json!({ + "grant_type":"urn:ietf:params:oauth:grant-type:jwt-bearer", + "assertion":response["access_token"],"client_id":"mcp" + }) + ); + } +} + +#[tokio::test] +async fn expired_grants_are_not_redeemed() { + let mut grant = exchange(&MockHttp::new(200, jag(&claims())), "") + .await + .unwrap(); + grant.expires_at = 0; + let resource_as = MockHttp::default(); + assert!(matches!( + grant.exchange(&resource_as).await, + Err(EmaError::InvalidRequest(_)) + )); + assert!(resource_as.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn bearer_responses_cannot_change_resource_or_widen_scope() { + let mut cases = vec![ + ("scope", json!("files.read"), true), + ("scope", json!("files.admin"), false), + ("scope", json!("files.read files.write"), false), + ("scope", json!("files.read files.read"), false), + ("resource", json!(RESOURCE), true), + ("resource", json!([RESOURCE]), true), + ("resource", json!([RESOURCE, "other"]), false), + ("resource", json!("https://mcp.example"), false), + ("expires_in", json!(0), false), + ("expires_in", Value::Null, true), + ("refresh_token", json!("unsupported"), false), + ("token_type", json!("N_A"), false), + ("token_type", json!("bearer"), true), + ("access_token", json!(" \t"), false), + ]; + cases.extend( + BAD_SCOPES + .iter() + .map(|scope| ("scope", json!(scope), false)), + ); + for (key, value, valid) in cases { + let mut token = bearer(); + if value.is_null() { + token.as_object_mut().unwrap().remove(key); + } else { + token[key] = value; + } + let result = redeem(&MockHttp::new(200, token)).await; + assert_eq!(result.is_ok(), valid, "{key}"); + if key == "expires_in" && valid { + assert_eq!(result.unwrap().expires_in, None); + } else if !valid { + assert!(matches!( + result, + Err(EmaError::InvalidResponse { + stage: ResourceServer, + .. + }) + )); + } + } +} + +#[tokio::test] +async fn resource_scope_narrowing_is_reported_without_logging_scope_values() { + for read in ["files.read", "private-scope-sentinel"] { + let requested = format!("{read} files.write"); + let mut claims = claims(); + claims["scope"] = json!(requested); + let grant = exchange(&MockHttp::new(200, jag(&claims)), &requested) + .await + .unwrap(); + let mut response = bearer(); + response["scope"] = json!(read); + let token = grant.exchange(&MockHttp::new(200, response)).await.unwrap(); + assert_eq!(token.scopes, HashSet::from([read.to_owned()])); + assert!(!format!("{token:?}").contains(read)); + } +} + +#[tokio::test] +async fn bearer_scope_may_be_omitted_but_not_added_to_an_unscoped_grant() { + for scope in [None, Some("files.read")] { + let mut claims = claims(); + claims.as_object_mut().unwrap().remove("scope"); + let grant = exchange(&MockHttp::new(200, jag(&claims)), "") + .await + .unwrap(); + let mut token = bearer(); + if let Some(scope) = scope { + token["scope"] = json!(scope); + } + let result = grant.exchange(&MockHttp::new(200, token)).await; + assert_eq!(result.is_ok(), scope.is_none()); + if let Ok(token) = result { + assert!(token.scopes.is_empty()); + } + } +} + +#[tokio::test] +async fn resource_errors_are_staged_and_never_reflect_credentials() { + const SECRET: &str = "secret-resource-error-sentinel"; + for (status, code) in [ + (400, "invalid_grant"), + (400, "insufficient_user_authentication"), + (400, "invalid_client"), + (302, SECRET), + (500, SECRET), + ] { + let http = MockHttp::new(status, json!({"error":code,"error_description":SECRET})); + let error = redeem(&http).await.unwrap_err(); + match code { + "invalid_grant" => assert_eq!(error, EmaError::InvalidGrant(ResourceServer)), + "insufficient_user_authentication" => assert_eq!( + error, + EmaError::InsufficientUserAuthentication(ResourceServer) + ), + _ => assert!( + matches!(error, EmaError::OAuthRejected {stage: ResourceServer, status: actual, ..} if actual == status) + ), + } + assert!(!format!("{error:?} {error}").contains(SECRET)); + assert!(std::error::Error::source(&error).is_none()); + assert_eq!(http.requests.lock().unwrap().len(), 1); + } + let mut oversized = bearer(); + oversized["ignored"] = json!("x".repeat(1024 * 1024)); + for (body, adapter_failure) in [ + (json!({"access_token":SECRET,"expires_in":SECRET}), false), + (oversized, false), + (Value::Null, true), + ] { + let http = MockHttp::new(200, body); + if adapter_failure { + *http.response.lock().unwrap() = Some(Err(SECRET.into())); + } + let error = redeem(&http).await.unwrap_err(); + assert!(!format!("{error:?} {error}").contains(SECRET)); + assert!(std::error::Error::source(&error).is_none()); + if adapter_failure { + assert_eq!(error, EmaError::RequestFailed(ResourceServer)); + } else { + assert!(matches!( + error, + EmaError::InvalidResponse { + stage: ResourceServer, + .. + } + )); + } + } +} + +#[tokio::test(start_paused = true)] +async fn both_exchanges_enforce_the_timeout_when_the_adapter_does_not() { + struct PendingHttp; + impl OAuthHttpClient for PendingHttp { + fn execute(&self, _: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { + Box::pin(std::future::pending()) + } + } + for stage in [Idp, ResourceServer] { + let ready = MockHttp::new(200, jag(&claims())); + let (idp, resource): (&dyn OAuthHttpClient, &dyn OAuthHttpClient) = match stage { + Idp => (&PendingHttp, &ready), + _ => (&ready, &PendingHttp), + }; + let refresh = RefreshToken::new("refresh-token".into()); + let request = EmaExchangeRequest::new( + EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp"), + EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp"), + RESOURCE, + &refresh, + ); + let start = tokio::time::Instant::now(); + let result = tokio::time::timeout(Duration::from_secs(31), request.exchange(idp, resource)) + .await + .expect("the SDK must enforce its own deadline"); + assert_eq!(result.unwrap_err(), EmaError::RequestFailed(stage)); + assert_eq!(start.elapsed(), Duration::from_secs(30)); + assert_eq!( + ready.requests.lock().unwrap().len(), + usize::from(stage == ResourceServer) + ); + } +} + +#[derive(Clone, Copy)] +enum AssertionBehavior { + Success, + Error, + Empty(&'static str), + Delay(Duration), + Pending, +} + +struct RecordingAssertionProvider { + calls: Mutex>, + behavior: AssertionBehavior, +} + +impl RecordingAssertionProvider { + fn new(behavior: AssertionBehavior) -> Arc { + Arc::new(Self { + calls: Mutex::new(Vec::new()), + behavior, + }) + } +} + +fn client_assertion(client_id: &str, issuer: &str, sequence: usize) -> String { + jwt( + json!({"alg":"ES256","typ":"client-authentication+jwt"}), + &json!({"iss":client_id,"sub":client_id,"aud":issuer, + "exp":4102444800_u64,"jti":format!("client-assertion-{sequence}")}), + ) +} + +#[async_trait::async_trait] +impl EmaClientAssertionProvider for RecordingAssertionProvider { + async fn create_assertion( + &self, + server: &EmaAuthorizationServer, + ) -> Result> { + let sequence = { + let mut calls = self.calls.lock().unwrap(); + calls.push(( + server.issuer.clone(), + server.token_endpoint.clone(), + server.client_id.clone(), + )); + calls.len() + }; + match self.behavior { + AssertionBehavior::Error => return Err("client-assertion-provider-secret".into()), + AssertionBehavior::Empty(value) => return Ok(EmaClientAssertion::new(value.into())), + AssertionBehavior::Delay(delay) => tokio::time::sleep(delay).await, + AssertionBehavior::Pending => std::future::pending::<()>().await, + AssertionBehavior::Success => {} + } + Ok(EmaClientAssertion::new(client_assertion( + &server.client_id, + &server.issuer, + sequence, + ))) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum AuthenticationMethod { + Public, + Basic, + Post, + Jwt, +} + +impl AuthenticationMethod { + fn configure( + self, + secret: &str, + provider: &Arc, + ) -> EmaClientAuthentication { + match self { + Self::Public => EmaClientAuthentication::None, + Self::Basic => { + EmaClientAuthentication::ClientSecretBasic(ClientSecret::new(secret.into())) + } + Self::Post => { + EmaClientAuthentication::ClientSecretPost(ClientSecret::new(secret.into())) + } + Self::Jwt => EmaClientAuthentication::JwtAssertion(provider.clone()), + } + } +} + +fn remove_client_authentication( + request: &OAuthHttpRequest, + method: AuthenticationMethod, + client_id: &str, + secret: &str, + encoded_basic: &str, + assertion: &str, +) -> BTreeMap { + let mut fields = authenticated_form(request); + let authorization = request.request.headers().get("authorization"); + if method == AuthenticationMethod::Basic { + let authorization = authorization.expect("Basic authentication must use a header"); + assert!(authorization.is_sensitive()); + let encoded = authorization + .to_str() + .unwrap() + .strip_prefix("Basic ") + .unwrap(); + assert_eq!(STANDARD.decode(encoded).unwrap(), encoded_basic.as_bytes()); + assert!(!fields.contains_key("client_id")); + } else { + assert!(authorization.is_none()); + assert_eq!(fields.remove("client_id").as_deref(), Some(client_id)); + } + assert_eq!( + fields.remove("client_secret").as_deref(), + (method == AuthenticationMethod::Post).then_some(secret) + ); + assert_eq!( + fields.remove("client_assertion_type").as_deref(), + (method == AuthenticationMethod::Jwt) + .then_some("urn:ietf:params:oauth:client-assertion-type:jwt-bearer") + ); + assert_eq!( + fields.remove("client_assertion").as_deref(), + (method == AuthenticationMethod::Jwt).then_some(assertion) + ); + fields +} + +#[tokio::test] +async fn client_authentication_is_endpoint_specific_and_never_mixed_with_grants() { + const METHODS: [AuthenticationMethod; 4] = [ + AuthenticationMethod::Public, + AuthenticationMethod::Basic, + AuthenticationMethod::Post, + AuthenticationMethod::Jwt, + ]; + // Colons, plus signs, spaces, percent signs, and non-ASCII bytes must be + // form-encoded individually before the HTTP Basic username/password join. + const IDP_CLIENT: &str = "idp:+ %é"; + const AS_CLIENT: &str = "mcp:+ %é"; + const IDP_SECRET: &str = "idp-secret:+ %é"; + const AS_SECRET: &str = "as-secret:+ %é"; + for idp_method in METHODS { + for resource_method in METHODS { + let provider = RecordingAssertionProvider::new(AssertionBehavior::Success); + let mut claims = claims(); + claims["client_id"] = json!(AS_CLIENT); + let response = jag(&claims); + let idp = MockHttp::new(200, response.clone()); + let resource = MockHttp::new(200, bearer()); + let refresh = RefreshToken::new("refresh-token".into()); + let token = EmaExchangeRequest::new( + EmaAuthorizationServer::new(IDP, IDP_TOKEN, IDP_CLIENT) + .with_client_authentication(idp_method.configure(IDP_SECRET, &provider)), + EmaAuthorizationServer::new(AS, AS_TOKEN, AS_CLIENT) + .with_client_authentication(resource_method.configure(AS_SECRET, &provider)), + RESOURCE, + &refresh, + ) + .exchange(&idp, &resource) + .await + .unwrap(); + assert_eq!(token.access_token.secret(), "resource-token"); + let idp_requests = idp.requests.lock().unwrap(); + let resource_requests = resource.requests.lock().unwrap(); + assert_eq!(idp_requests.len(), 1); + assert_eq!(resource_requests.len(), 1); + assert_eq!(idp_requests[0].request.uri(), IDP_TOKEN); + assert_eq!(resource_requests[0].request.uri(), AS_TOKEN); + let idp_fields = remove_client_authentication( + &idp_requests[0], + idp_method, + IDP_CLIENT, + IDP_SECRET, + "idp%3A%2B+%25%C3%A9:idp-secret%3A%2B+%25%C3%A9", + &client_assertion(IDP_CLIENT, IDP, 1), + ); + assert_eq!( + serde_json::to_value(idp_fields).unwrap(), + json!({ + "grant_type":"urn:ietf:params:oauth:grant-type:token-exchange", + "requested_token_type":ID_JAG_TOKEN_TYPE,"subject_token":"refresh-token", + "subject_token_type":"urn:ietf:params:oauth:token-type:refresh_token", + "audience":AS,"resource":RESOURCE + }) + ); + let resource_fields = remove_client_authentication( + &resource_requests[0], + resource_method, + AS_CLIENT, + AS_SECRET, + "mcp%3A%2B+%25%C3%A9:as-secret%3A%2B+%25%C3%A9", + &client_assertion( + AS_CLIENT, + AS, + 1 + usize::from(idp_method == AuthenticationMethod::Jwt), + ), + ); + assert_eq!( + serde_json::to_value(resource_fields).unwrap(), + json!({ + "grant_type":"urn:ietf:params:oauth:grant-type:jwt-bearer", + "assertion":response["access_token"] + }) + ); + let mut expected_calls = Vec::new(); + if idp_method == AuthenticationMethod::Jwt { + expected_calls.push((IDP.into(), IDP_TOKEN.into(), IDP_CLIENT.into())); + } + if resource_method == AuthenticationMethod::Jwt { + expected_calls.push((AS.into(), AS_TOKEN.into(), AS_CLIENT.into())); + } + assert_eq!(*provider.calls.lock().unwrap(), expected_calls); + } + } +} + +#[tokio::test(start_paused = true)] +async fn client_assertions_are_fresh_for_every_request_and_delayed_redemption() { + let provider = RecordingAssertionProvider::new(AssertionBehavior::Success); + let idp_server = EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp") + .with_client_authentication(EmaClientAuthentication::JwtAssertion(provider.clone())); + let resource_server = EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp") + .with_client_authentication(EmaClientAuthentication::JwtAssertion(provider.clone())); + let refresh = RefreshToken::new("refresh-token".into()); + let mut assertions = HashSet::new(); + for attempt in 0..2 { + let idp = MockHttp::new(200, jag(&claims())); + let resource = MockHttp::new(200, bearer()); + let grant = EmaExchangeRequest::new( + idp_server.clone(), + resource_server.clone(), + RESOURCE, + &refresh, + ) + .exchange_id_jag(&idp) + .await + .unwrap(); + assert_eq!(provider.calls.lock().unwrap().len(), attempt * 2 + 1); + tokio::time::sleep(Duration::from_secs(60)).await; + assert_eq!(provider.calls.lock().unwrap().len(), attempt * 2 + 1); + grant.exchange(&resource).await.unwrap(); + assert_eq!(provider.calls.lock().unwrap().len(), attempt * 2 + 2); + for http in [&idp, &resource] { + let requests = http.requests.lock().unwrap(); + let mut fields = form(&requests[0]); + assert!(assertions.insert(fields.remove("client_assertion").unwrap())); + } + } + assert_eq!(assertions.len(), 4); +} + +#[tokio::test] +async fn invalid_static_client_secrets_fail_before_any_http_or_signing() { + for stage in [Idp, ResourceServer] { + for method in [AuthenticationMethod::Basic, AuthenticationMethod::Post] { + for secret in ["", " \t"] { + let provider = RecordingAssertionProvider::new(AssertionBehavior::Success); + let valid = EmaClientAuthentication::JwtAssertion(provider.clone()); + let invalid = method.configure(secret, &provider); + let (idp_auth, resource_auth) = match stage { + Idp => (invalid, valid), + _ => (valid, invalid), + }; + let idp = MockHttp::default(); + let resource = MockHttp::default(); + let error = EmaExchangeRequest::new( + EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp") + .with_client_authentication(idp_auth), + EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp") + .with_client_authentication(resource_auth), + RESOURCE, + &RefreshToken::new("refresh-token".into()), + ) + .exchange(&idp, &resource) + .await + .unwrap_err(); + assert!(matches!(error, EmaError::InvalidRequest(_))); + assert!(idp.requests.lock().unwrap().is_empty()); + assert!(resource.requests.lock().unwrap().is_empty()); + assert!(provider.calls.lock().unwrap().is_empty()); + } + } + } +} + +#[tokio::test] +async fn a_grant_that_expires_while_signing_is_never_sent_to_the_resource_server() { + use std::sync::atomic::{AtomicU64, Ordering}; + + let provider = RecordingAssertionProvider::new(AssertionBehavior::Success); + let idp = MockHttp::new(200, jag(&claims())); + let resource = MockHttp::default(); + let mut grant = EmaExchangeRequest::new( + EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp"), + EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp") + .with_client_authentication(EmaClientAuthentication::JwtAssertion(provider.clone())), + RESOURCE, + &RefreshToken::new("refresh-token".into()), + ) + .exchange_id_jag(&idp) + .await + .unwrap(); + grant.expires_at = 100; + // The clock crosses expiration between the checks before and after signing. + let now = AtomicU64::new(99); + let error = grant + .exchange_with_clock(&resource, || Ok(now.fetch_add(1, Ordering::Relaxed))) + .await + .unwrap_err(); + assert!(matches!(error, EmaError::InvalidRequest(_))); + assert_eq!(provider.calls.lock().unwrap().len(), 1); + assert!(resource.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn assertion_provider_failures_and_empty_results_never_reach_http_or_escape_errors() { + for stage in [Idp, ResourceServer] { + for behavior in [ + AssertionBehavior::Error, + AssertionBehavior::Empty(""), + AssertionBehavior::Empty(" \t"), + ] { + let provider = RecordingAssertionProvider::new(behavior); + let auth = EmaClientAuthentication::JwtAssertion(provider.clone()); + let (idp_auth, resource_auth) = match stage { + Idp => (auth, EmaClientAuthentication::None), + _ => (EmaClientAuthentication::None, auth), + }; + let idp = MockHttp::new(200, jag(&claims())); + let resource = MockHttp::default(); + let error = EmaExchangeRequest::new( + EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp") + .with_client_authentication(idp_auth), + EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp") + .with_client_authentication(resource_auth), + RESOURCE, + &RefreshToken::new("refresh-token".into()), + ) + .exchange(&idp, &resource) + .await + .unwrap_err(); + if matches!(behavior, AssertionBehavior::Error) { + assert_eq!(error, EmaError::RequestFailed(stage)); + } else { + assert!(matches!(error, EmaError::InvalidRequest(_))); + } + assert!(!format!("{error:?} {error}").contains("client-assertion-provider-secret")); + assert!(std::error::Error::source(&error).is_none()); + assert_eq!(provider.calls.lock().unwrap().len(), 1); + assert_eq!( + idp.requests.lock().unwrap().len(), + usize::from(stage == ResourceServer) + ); + assert!(resource.requests.lock().unwrap().is_empty()); + } + } +} + +#[tokio::test(start_paused = true)] +async fn signing_and_http_share_one_deadline_at_both_endpoints() { + struct DelayedHttp(Mutex>); + impl OAuthHttpClient for DelayedHttp { + fn execute(&self, request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { + self.0.lock().unwrap().push(request); + Box::pin(async { + tokio::time::sleep(Duration::from_secs(10)).await; + Err("http-adapter-secret".into()) + }) + } + } + for stage in [Idp, ResourceServer] { + for behavior in [ + AssertionBehavior::Pending, + AssertionBehavior::Delay(Duration::from_secs(25)), + ] { + let provider = RecordingAssertionProvider::new(behavior); + let auth = EmaClientAuthentication::JwtAssertion(provider.clone()); + let (idp_auth, resource_auth) = match stage { + Idp => (auth, EmaClientAuthentication::None), + _ => (EmaClientAuthentication::None, auth), + }; + let ready = MockHttp::new(200, jag(&claims())); + let delayed = DelayedHttp(Mutex::new(Vec::new())); + let (idp, resource): (&dyn OAuthHttpClient, &dyn OAuthHttpClient) = match stage { + Idp => (&delayed, &ready), + _ => (&ready, &delayed), + }; + let refresh = RefreshToken::new("refresh-token".into()); + let request = EmaExchangeRequest::new( + EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp") + .with_client_authentication(idp_auth), + EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp") + .with_client_authentication(resource_auth), + RESOURCE, + &refresh, + ); + let start = tokio::time::Instant::now(); + let result = + tokio::time::timeout(Duration::from_secs(31), request.exchange(idp, resource)) + .await + .expect("signing must share the SDK's token request deadline"); + assert_eq!(result.unwrap_err(), EmaError::RequestFailed(stage)); + assert_eq!(start.elapsed(), Duration::from_secs(30)); + assert_eq!(provider.calls.lock().unwrap().len(), 1); + assert_eq!( + delayed.0.lock().unwrap().len(), + usize::from(matches!(behavior, AssertionBehavior::Delay(_))) + ); + assert_eq!( + ready.requests.lock().unwrap().len(), + usize::from(stage == ResourceServer) + ); + } + } +} + +#[tokio::test] +async fn rejected_client_authentication_never_falls_back_or_retries() { + for stage in [Idp, ResourceServer] { + for method in [ + AuthenticationMethod::Basic, + AuthenticationMethod::Post, + AuthenticationMethod::Jwt, + ] { + let provider = RecordingAssertionProvider::new(AssertionBehavior::Success); + let failure = json!({"error":"invalid_client","error_description":"client-authentication-secret"}); + let idp = if stage == Idp { + MockHttp::new(401, failure.clone()) + } else { + MockHttp::new(200, jag(&claims())) + }; + let resource = MockHttp::new(401, failure); + let error = EmaExchangeRequest::new( + EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp") + .with_client_authentication(method.configure("idp-client-secret", &provider)), + EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp") + .with_client_authentication(method.configure("as-client-secret", &provider)), + RESOURCE, + &RefreshToken::new("refresh-token".into()), + ) + .exchange(&idp, &resource) + .await + .unwrap_err(); + assert_eq!( + error, + EmaError::OAuthRejected { + stage, + status: 401, + code: "invalid_client" + } + ); + assert!(!format!("{error:?} {error}").contains("client-authentication-secret")); + assert_eq!(idp.requests.lock().unwrap().len(), 1); + assert_eq!( + resource.requests.lock().unwrap().len(), + usize::from(stage == ResourceServer) + ); + let expected_calls = if method == AuthenticationMethod::Jwt { + 1 + usize::from(stage == ResourceServer) + } else { + 0 + }; + assert_eq!(provider.calls.lock().unwrap().len(), expected_calls); + } + } +} + +#[test] +fn client_authentication_debug_output_redacts_credentials_and_provider_state() { + const SECRET: &str = "client-authentication-secret-sentinel"; + let provider = RecordingAssertionProvider::new(AssertionBehavior::Empty(SECRET)); + let assertion = EmaClientAssertion::new(SECRET.into()); + assert!(!format!("{assertion:?}").contains(SECRET)); + for authentication in [ + EmaClientAuthentication::ClientSecretBasic(ClientSecret::new(SECRET.into())), + EmaClientAuthentication::ClientSecretPost(ClientSecret::new(SECRET.into())), + EmaClientAuthentication::JwtAssertion(provider.clone()), + ] { + assert!(!format!("{authentication:?}").contains(SECRET)); + let server = EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp") + .with_client_authentication(authentication); + let refresh = RefreshToken::new("refresh-token".into()); + let request = EmaExchangeRequest::new(server.clone(), server.clone(), RESOURCE, &refresh); + assert!(!format!("{server:?} {request:?}").contains(SECRET)); + assert!(!format!("{server:?} {request:?}").contains("private-query")); + } +} diff --git a/crates/rmcp/src/transport/child_process.rs b/crates/rmcp/src/transport/child_process.rs index ebb6cc9..6e19a0c 100644 --- a/crates/rmcp/src/transport/child_process.rs +++ b/crates/rmcp/src/transport/child_process.rs @@ -2,10 +2,7 @@ use std::process::Stdio; use futures::future::Future; use process_wrap::tokio::{ChildWrapper, CommandWrap}; -use tokio::{ - io::AsyncRead, - process::{ChildStderr, ChildStdin, ChildStdout}, -}; +use tokio::process::{ChildStderr, ChildStdin, ChildStdout}; use super::{RxJsonRpcMessage, Transport, TxJsonRpcMessage, async_rw::AsyncRwTransport}; use crate::RoleClient; @@ -59,32 +56,6 @@ impl Drop for ChildWithCleanup { } } -// we hold the child process with stdout, for it's easier to implement AsyncRead -pin_project_lite::pin_project! { - pub struct TokioChildProcessOut { - child: ChildWithCleanup, - #[pin] - child_stdout: ChildStdout, - } -} - -impl TokioChildProcessOut { - /// Get the process ID of the child process. - pub fn id(&self) -> Option { - self.child.inner.as_ref()?.id() - } -} - -impl AsyncRead for TokioChildProcessOut { - fn poll_read( - self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - buf: &mut tokio::io::ReadBuf<'_>, - ) -> std::task::Poll> { - self.project().child_stdout.poll_read(cx, buf) - } -} - impl TokioChildProcess { /// Convenience: spawn with default `piped` stdio pub fn new(command: impl Into) -> std::io::Result { @@ -139,15 +110,6 @@ impl TokioChildProcess { pub fn into_inner(mut self) -> Option> { self.child.inner.take() } - - /// Split this helper into a reader (stdout) and writer (stdin). - #[deprecated( - since = "0.5.0", - note = "use the Transport trait implementation instead" - )] - pub fn split(self) -> (TokioChildProcessOut, ChildStdin) { - unimplemented!("This method is deprecated, use the Transport trait implementation instead"); - } } /// Builder for `TokioChildProcess` allowing custom `Stdio` configuration. diff --git a/crates/rmcp/src/transport/common.rs b/crates/rmcp/src/transport/common.rs index 3691602..8cf00c0 100644 --- a/crates/rmcp/src/transport/common.rs +++ b/crates/rmcp/src/transport/common.rs @@ -3,6 +3,9 @@ pub mod server_side_http; pub mod http_header; +#[cfg(any(feature = "client-side-sse", feature = "server-side-http"))] +pub mod mcp_headers; + #[cfg(feature = "__reqwest")] mod reqwest; diff --git a/crates/rmcp/src/transport/common/auth/streamable_http_client.rs b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs index 47f08f1..47b6caf 100644 --- a/crates/rmcp/src/transport/common/auth/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs @@ -1,11 +1,76 @@ use std::collections::HashMap; use http::{HeaderName, HeaderValue}; +use tracing::debug; use crate::transport::{ - auth::AuthClient, + auth::{AuthClient, AuthError}, streamable_http_client::{StreamableHttpClient, StreamableHttpError}, }; + +impl AuthClient +where + C: StreamableHttpClient + Send + Sync, +{ + /// Run `call` with a token when one is available, reacting to the + /// server's auth verdict instead of requiring credentials up front: + /// + /// - no usable credentials → the request goes out unauthenticated, and a + /// 401 propagates as [`StreamableHttpError::AuthRequired`] carrying the + /// `WWW-Authenticate` challenge for the caller to authorize with; + /// - a token the server rejects (e.g. revoked) → one silent refresh, one + /// retry, then the challenge propagates; + /// - a refresh that fails for any other reason (credential store, network, + /// provider) → that error propagates so the caller can retry instead of + /// being sent through a new authorization. + async fn call_reacting_to_challenges( + &self, + auth_token: Option, + call: F, + ) -> Result> + where + F: Fn(Option) -> Fut, + Fut: Future>>, + { + // Missing credentials are not an error in the reactive model: the + // request goes out unauthenticated and the server's 401 challenge + // drives authorization. + let auth_token = match auth_token { + None => match self.get_access_token().await { + Ok(token) => Some(token), + Err(AuthError::AuthorizationRequired) => None, + Err(error) => return Err(error.into()), + }, + token => token, + }; + match call(auth_token.clone()).await { + Err(StreamableHttpError::AuthRequired(challenge)) => { + // One silent recovery attempt: refresh the rejected token and + // retry only when the refresh actually produced a new one. + let Some(sent_token) = auth_token else { + return Err(StreamableHttpError::AuthRequired(challenge)); + }; + let refreshed = { + let manager = self.auth_manager.lock().await; + manager.try_refresh_or_reauth().await + }; + match refreshed { + Ok(fresh_token) if fresh_token != sent_token => call(Some(fresh_token)).await, + Ok(_) => Err(StreamableHttpError::AuthRequired(challenge)), + // `try_refresh_or_reauth` already reports the cases that need a + // new authorization; anything else is retryable or infrastructural. + Err(AuthError::AuthorizationRequired) => { + debug!("token refresh after server rejection requires authorization"); + Err(StreamableHttpError::AuthRequired(challenge)) + } + Err(error) => Err(error.into()), + } + } + result => result, + } + } +} + impl StreamableHttpClient for AuthClient where C: StreamableHttpClient + Send + Sync, @@ -16,35 +81,79 @@ where &self, uri: std::sync::Arc, session_id: std::sync::Arc, - mut auth_token: Option, + auth_token: Option, custom_headers: HashMap, ) -> Result<(), crate::transport::streamable_http_client::StreamableHttpError> { - if auth_token.is_none() { - auth_token = Some(self.get_access_token().await?); - } - self.http_client - .delete_session(uri, session_id, auth_token, custom_headers) - .await + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let session_id = session_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .delete_session(uri, session_id, token, custom_headers) + .await + } + }) + .await } async fn get_stream( &self, uri: std::sync::Arc, - session_id: std::sync::Arc, + session_id: Option>, last_event_id: Option, - mut auth_token: Option, + auth_token: Option, custom_headers: HashMap, ) -> Result< futures::stream::BoxStream<'static, Result>, crate::transport::streamable_http_client::StreamableHttpError, > { - if auth_token.is_none() { - auth_token = Some(self.get_access_token().await?); - } - self.http_client - .get_stream(uri, session_id, last_event_id, auth_token, custom_headers) - .await + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let session_id = session_id.clone(); + let last_event_id = last_event_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .get_stream(uri, session_id, last_event_id, token, custom_headers) + .await + } + }) + .await + } + + async fn get_stream_with_max_sse_event_size( + &self, + uri: std::sync::Arc, + session_id: Option>, + last_event_id: Option, + auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, + ) -> Result< + futures::stream::BoxStream<'static, Result>, + crate::transport::streamable_http_client::StreamableHttpError, + > { + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let session_id = session_id.clone(); + let last_event_id = last_event_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + token, + custom_headers, + max_sse_event_size, + ) + .await + } + }) + .await } async fn post_message( @@ -52,17 +161,220 @@ where uri: std::sync::Arc, message: crate::model::ClientJsonRpcMessage, session_id: Option>, - mut auth_token: Option, + auth_token: Option, custom_headers: HashMap, ) -> Result< crate::transport::streamable_http_client::StreamableHttpPostResponse, StreamableHttpError, > { - if auth_token.is_none() { - auth_token = Some(self.get_access_token().await?); + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let message = message.clone(); + let session_id = session_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .post_message(uri, message, session_id, token, custom_headers) + .await + } + }) + .await + } + + async fn post_message_with_max_sse_event_size( + &self, + uri: std::sync::Arc, + message: crate::model::ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, + ) -> Result< + crate::transport::streamable_http_client::StreamableHttpPostResponse, + StreamableHttpError, + > { + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let message = message.clone(); + let session_id = session_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .post_message_with_max_sse_event_size( + uri, + message, + session_id, + token, + custom_headers, + max_sse_event_size, + ) + .await + } + }) + .await + } +} + +#[cfg(all(test, feature = "transport-streamable-http-client-reqwest"))] +mod tests { + use std::sync::Arc; + + use oauth2::{AccessToken, RefreshToken, basic::BasicTokenType}; + + use super::*; + use crate::transport::{ + auth::{ + AuthorizationManager, AuthorizationMetadata, CredentialRefreshGuard, CredentialStore, + InMemoryCredentialStore, OAuthHttpClient, OAuthHttpClientFuture, OAuthHttpRequest, + OAuthTokenResponse, StoredCredentials, VendorExtraTokenFields, + }, + streamable_http_client::AuthRequiredError, + }; + + struct UnavailableStore; + + #[async_trait::async_trait] + impl CredentialStore for UnavailableStore { + async fn load(&self) -> Result, AuthError> { + unreachable!("guard failure must stop the credential load") } - self.http_client - .post_message(uri, message, session_id, auth_token, custom_headers) + + async fn save(&self, _: StoredCredentials) -> Result<(), AuthError> { + unreachable!("guard failure must stop the credential save") + } + + async fn clear(&self) -> Result<(), AuthError> { + unreachable!("refresh must not clear credentials") + } + + async fn acquire_refresh_guard(&self) -> Result, AuthError> { + Err(AuthError::CredentialStoreError("guard unavailable".into())) + } + } + + #[tokio::test] + async fn reactive_refresh_preserves_credential_store_failure() { + let mut manager = AuthorizationManager::new("https://mcp.example.com/mcp") + .await + .unwrap(); + manager.set_metadata(AuthorizationMetadata { + authorization_endpoint: "https://auth.example.com/authorize".into(), + token_endpoint: "https://auth.example.com/token".into(), + ..Default::default() + }); + manager.configure_client_id("client").unwrap(); + manager.set_credential_store(UnavailableStore); + let client = AuthClient::new(reqwest::Client::new(), manager); + + let error = client + .call_reacting_to_challenges(Some("old-token".into()), |_| async { + Err::<(), _>(StreamableHttpError::AuthRequired(AuthRequiredError::new( + "Bearer".into(), + ))) + }) + .await + .unwrap_err(); + + assert!(matches!(error, + StreamableHttpError::Auth(AuthError::CredentialStoreError(message)) + if message == "guard unavailable")); + } + + struct UnreachableTokenEndpoint; + + impl OAuthHttpClient for UnreachableTokenEndpoint { + fn execute(&self, _: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { + Box::pin(async { Err("token endpoint unreachable".into()) }) + } + } + + struct RejectingTokenEndpoint; + + impl OAuthHttpClient for RejectingTokenEndpoint { + fn execute(&self, _: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { + Box::pin(async { + Ok(oauth2::http::Response::builder() + .status(400) + .header("content-type", "application/json") + .body(br#"{"error":"invalid_grant"}"#.to_vec()) + .unwrap()) + }) + } + } + + /// A manager holding a refresh token the given token endpoint will answer for. + async fn manager_with_stored_refresh_token( + token_endpoint: Arc, + ) -> AuthorizationManager { + let mut manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + token_endpoint, + ) + .await + .unwrap(); + manager.set_metadata(AuthorizationMetadata { + authorization_endpoint: "https://auth.example.com/authorize".into(), + token_endpoint: "https://auth.example.com/token".into(), + ..Default::default() + }); + manager.configure_client_id("client").unwrap(); + + let mut token_response = OAuthTokenResponse::new( + AccessToken::new("old-token".into()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + token_response.set_refresh_token(Some(RefreshToken::new("stored-refresh".into()))); + let store = InMemoryCredentialStore::new(); + store + .save(StoredCredentials::new( + "client".into(), + Some(token_response), + vec![], + None, + )) + .await + .unwrap(); + manager.set_credential_store(store); + manager + } + + /// Drive one call whose server answer is a 401 challenge. + async fn challenge_once(manager: AuthorizationManager) -> StreamableHttpError { + AuthClient::new(reqwest::Client::new(), manager) + .call_reacting_to_challenges(Some("old-token".into()), |_| async { + Err::<(), _>(StreamableHttpError::AuthRequired(AuthRequiredError::new( + "Bearer".into(), + ))) + }) .await + .unwrap_err() + } + + #[tokio::test] + async fn reactive_refresh_propagates_retryable_refresh_failure() { + let manager = manager_with_stored_refresh_token(Arc::new(UnreachableTokenEndpoint)).await; + + let error = challenge_once(manager).await; + + assert!( + matches!( + error, + StreamableHttpError::Auth(AuthError::TokenRefreshFailed(_)) + ), + "a retryable refresh failure must reach the caller instead of asking for a new authorization, got: {error:?}" + ); + } + + #[tokio::test] + async fn reactive_refresh_reports_a_rejected_refresh_token_as_a_challenge() { + let manager = manager_with_stored_refresh_token(Arc::new(RejectingTokenEndpoint)).await; + + let error = challenge_once(manager).await; + + assert!( + matches!(error, StreamableHttpError::AuthRequired(_)), + "a definitively rejected refresh token must surface the challenge, got: {error:?}" + ); } } diff --git a/crates/rmcp/src/transport/common/client_side_sse.rs b/crates/rmcp/src/transport/common/client_side_sse.rs index fc9e15e..2ee37aa 100644 --- a/crates/rmcp/src/transport/common/client_side_sse.rs +++ b/crates/rmcp/src/transport/common/client_side_sse.rs @@ -5,13 +5,166 @@ use std::{ time::Duration, }; -use futures::{Stream, stream::BoxStream}; -use sse_stream::{Error as SseError, Sse}; +use bytes::Bytes; +use futures::{Stream, StreamExt, stream::BoxStream}; +use sse_stream::{Error as SseError, Sse, SseStream}; +use thiserror::Error; use crate::model::ServerJsonRpcMessage; pub type BoxedSseResponse = BoxStream<'static, Result>; +/// Maximum raw size of one SSE event accepted from a remote server. +pub(crate) const DEFAULT_MAX_SSE_EVENT_SIZE: usize = 16 * 1024 * 1024; + +#[derive(Debug, Error)] +enum BoundedSseStreamError { + #[error(transparent)] + Source(Box), + #[error("SSE event exceeded the maximum size of {max_size} bytes")] + EventTooLarge { max_size: usize }, +} + +#[derive(Debug)] +struct SseEventSizeLimiter { + max_size: usize, + retained_size: usize, + line_size: usize, + line_is_comment: bool, + previous_was_cr: bool, +} + +impl SseEventSizeLimiter { + fn new(max_size: usize) -> Self { + Self { + max_size, + retained_size: 0, + line_size: 0, + line_is_comment: false, + previous_was_cr: false, + } + } + + fn observe(&mut self, chunk: &[u8]) -> Result<(), ()> { + for &byte in chunk { + if self.previous_was_cr { + self.previous_was_cr = false; + if byte == b'\n' { + continue; + } + } + + match byte { + b'\r' => { + self.finish_line()?; + self.previous_was_cr = true; + } + b'\n' => self.finish_line()?, + _ => { + if self.line_size == 0 { + self.line_is_comment = byte == b':'; + } + self.line_size = self.line_size.saturating_add(1); + self.check_limit()?; + } + } + } + Ok(()) + } + + fn finish_line(&mut self) -> Result<(), ()> { + if self.line_size == 0 { + self.retained_size = 0; + } else if !self.line_is_comment { + // The SSE parser inserts a newline when joining multiple data fields. + self.retained_size = self + .retained_size + .saturating_add(self.line_size) + .saturating_add(1); + } + self.line_size = 0; + self.line_is_comment = false; + self.check_limit() + } + + fn check_limit(&self) -> Result<(), ()> { + if self.retained_size.saturating_add(self.line_size) > self.max_size { + Err(()) + } else { + Ok(()) + } + } +} + +pin_project_lite::pin_project! { + struct BoundedSseByteStream { + #[pin] + inner: S, + limiter: SseEventSizeLimiter, + failed: bool, + } +} + +impl Stream for BoundedSseByteStream +where + S: Stream>, + E: std::error::Error + Send + Sync + 'static, +{ + type Item = Result; + + fn poll_next( + self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll> { + let mut this = self.project(); + if *this.failed { + return Poll::Ready(None); + } + + match ready!(this.inner.as_mut().poll_next(cx)) { + Some(Ok(chunk)) => { + if this.limiter.observe(&chunk).is_err() { + *this.failed = true; + Poll::Ready(Some(Err(BoundedSseStreamError::EventTooLarge { + max_size: this.limiter.max_size, + }))) + } else { + Poll::Ready(Some(Ok(chunk))) + } + } + Some(Err(error)) => { + *this.failed = true; + Poll::Ready(Some(Err(BoundedSseStreamError::Source(Box::new(error))))) + } + None => Poll::Ready(None), + } + } +} + +pub(crate) fn bounded_sse_stream(stream: S, max_event_size: usize) -> BoxedSseResponse +where + S: Stream> + Send + 'static, + E: std::error::Error + Send + Sync + 'static, +{ + let stream = BoundedSseByteStream { + inner: stream, + limiter: SseEventSizeLimiter::new(max_event_size), + failed: false, + }; + SseStream::from_bytes_stream(stream).boxed() +} + +fn is_event_too_large_error(error: &SseError) -> bool { + matches!( + error, + SseError::Body(error) + if matches!( + error.downcast_ref::(), + Some(BoundedSseStreamError::EventTooLarge { .. }) + ) + ) +} + pub trait SseRetryPolicy: std::fmt::Debug + Send + Sync { fn retry(&self, current_times: usize) -> Option; } @@ -25,10 +178,10 @@ pub struct FixedInterval { impl SseRetryPolicy for FixedInterval { fn retry(&self, current_times: usize) -> Option { - if let Some(max_times) = self.max_times { - if current_times >= max_times { - return None; - } + if let Some(max_times) = self.max_times + && current_times >= max_times + { + return None; } Some(self.duration) } @@ -52,6 +205,11 @@ impl Default for FixedInterval { pub struct ExponentialBackoff { pub max_times: Option, pub base_duration: Duration, + /// Optional upper bound on a single reconnect delay. `None` (the default) preserves the + /// pre-existing unbounded doubling behavior. Once the multiplier saturates near the bit + /// width that can still produce very long sleeps, so callers that need the client to + /// actually reconnect can set `Some(...)` to clamp the delay. + pub max_delay: Option, } impl ExponentialBackoff { @@ -63,18 +221,28 @@ impl Default for ExponentialBackoff { Self { max_times: None, base_duration: Self::DEFAULT_DURATION, + max_delay: None, } } } impl SseRetryPolicy for ExponentialBackoff { fn retry(&self, current_times: usize) -> Option { - if let Some(max_times) = self.max_times { - if current_times >= max_times { - return None; - } + if let Some(max_times) = self.max_times + && current_times >= max_times + { + return None; } - Some(self.base_duration * (2u32.pow(current_times as u32))) + // `current_times` is unbounded when `max_times` is unset, so the exponent can reach + // the bit width. Saturate the multiplier at `u32::MAX` and use saturating multiplication + // for the base duration so the delay stays monotonic and panic-free instead of an + // overflow panic (debug) or a wrapped-to-zero backoff (release). + let multiplier = 2u32.saturating_pow(current_times as u32); + let delay = self.base_duration.saturating_mul(multiplier); + Some(match self.max_delay { + Some(max_delay) => delay.min(max_delay), + None => delay, + }) } } @@ -124,6 +292,10 @@ pub(crate) trait SseStreamReconnect { tracing::warn!("sse stream error: {error}"); } } + fn map_fatal_stream_error(&mut self, error: SseError) -> Option { + tracing::warn!("fatal sse stream error: {error}"); + None + } } pin_project_lite::pin_project! { @@ -131,6 +303,7 @@ pin_project_lite::pin_project! { where R: SseStreamReconnect { retry_policy: Arc, + reconnect_only_after_event_id: bool, last_event_id: Option, server_retry_interval: Option, connector: R, @@ -147,6 +320,22 @@ impl SseAutoReconnectStream { ) -> Self { Self { retry_policy, + reconnect_only_after_event_id: false, + last_event_id: None, + server_retry_interval: None, + connector, + state: SseAutoReconnectStreamState::Connected { stream }, + } + } + + pub fn new_after_event_id( + stream: BoxedSseResponse, + connector: R, + retry_policy: Arc, + ) -> Self { + Self { + retry_policy, + reconnect_only_after_event_id: true, last_event_id: None, server_retry_interval: None, connector, @@ -160,6 +349,7 @@ impl SseAutoReconnectStream> { pub(crate) fn never_reconnect(stream: BoxedSseResponse, error_when_reconnect: E) -> Self { Self { retry_policy: Arc::new(NeverRetry), + reconnect_only_after_event_id: false, last_event_id: None, server_retry_interval: None, connector: NeverReconnect { @@ -201,129 +391,474 @@ where mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> Poll> { - let mut this = self.as_mut().project(); - // let this_state = this.state.as_mut().project() - let state = this.state.as_mut().project(); - let next_state = match state { - SseAutoReconnectStreamStateProj::Connected { stream } => { - match ready!(stream.poll_next(cx)) { - Some(Ok(sse)) => { - if let Some(new_server_retry) = sse.retry { - *this.server_retry_interval = - Some(Duration::from_millis(new_server_retry)); - } - if let Some(ref event_id) = sse.id { - *this.last_event_id = Some(event_id.clone()); - } - // Only treat blank/`message` events as JSON-RPC payloads. - // Other control frames (endpoint, ping, etc.) are passed to - // the reconnection handler. - let is_message_event = - matches!(sse.event.as_deref(), None | Some("") | Some("message")); - if !is_message_event { - match this.connector.handle_control_event(&sse) { - Ok(()) => return self.poll_next(cx), - Err(e) => { - this.state.set(SseAutoReconnectStreamState::Terminated); - return Poll::Ready(Some(Err(e))); + loop { + let mut this = self.as_mut().project(); + // let this_state = this.state.as_mut().project() + let state = this.state.as_mut().project(); + let next_state = match state { + SseAutoReconnectStreamStateProj::Connected { stream } => { + match ready!(stream.poll_next(cx)) { + Some(Ok(sse)) => { + if let Some(new_server_retry) = sse.retry { + *this.server_retry_interval = + Some(Duration::from_millis(new_server_retry)); + } + if let Some(ref event_id) = sse.id { + *this.last_event_id = Some(event_id.clone()); + } + // Only treat blank/`message` events as JSON-RPC payloads. + // Other control frames (endpoint, ping, etc.) are passed to + // the reconnection handler. + let is_message_event = + matches!(sse.event.as_deref(), None | Some("") | Some("message")); + if !is_message_event { + match this.connector.handle_control_event(&sse) { + Ok(()) => continue, + Err(e) => { + this.state.set(SseAutoReconnectStreamState::Terminated); + return Poll::Ready(Some(Err(e))); + } } } + if let Some(data) = sse.data { + match serde_json::from_str::(&data) { + Err(e) => { + // Downgrade to debug to avoid noisy logs when servers emit + // non-JSON payloads as message frames. Include last_event_id + // to aid troubleshooting while keeping default behaviour. + let last_id = this.last_event_id.as_deref().unwrap_or(""); + tracing::debug!(last_event_id=%last_id, "failed to deserialize server message: {e}"); + continue; + } + Ok(message) => { + return Poll::Ready(Some(Ok(message))); + } + }; + } else { + continue; + } } - if let Some(data) = sse.data { - match serde_json::from_str::(&data) { - Err(e) => { - // Downgrade to debug to avoid noisy logs when servers emit - // non-JSON payloads as message frames. Include last_event_id - // to aid troubleshooting while keeping default behaviour. - let last_id = this.last_event_id.as_deref().unwrap_or(""); - tracing::debug!(last_event_id=%last_id, "failed to deserialize server message: {e}"); - return self.poll_next(cx); - } - Ok(message) => { - return Poll::Ready(Some(Ok(message))); - } - }; - } else { - return self.poll_next(cx); - } - } - Some(Err(e)) => { - this.connector - .handle_stream_error(&e, this.last_event_id.as_deref()); - let retrying = this - .connector - .retry_connection(this.last_event_id.as_deref()); - SseAutoReconnectStreamState::Retrying { - retry_times: 0, - retrying, - } - } - None => { - // Per SEP-1699, a graceful stream close is - // reconnectable. If the server sent a `retry` field - // we MUST wait that long before reconnecting. - let interval = this - .server_retry_interval - .take() - .or_else(|| this.retry_policy.retry(0)); - if let Some(interval) = interval { - tracing::debug!(?interval, "sse stream ended gracefully, reconnecting"); - SseAutoReconnectStreamState::WaitingNextRetry { - sleep: tokio::time::sleep(interval), + Some(Err(e)) => { + if is_event_too_large_error(&e) { + this.state.set(SseAutoReconnectStreamState::Terminated); + return Poll::Ready( + this.connector.map_fatal_stream_error(e).map(Err), + ); + } + if *this.reconnect_only_after_event_id && this.last_event_id.is_none() { + this.state.set(SseAutoReconnectStreamState::Terminated); + return Poll::Ready( + this.connector.map_fatal_stream_error(e).map(Err), + ); + } + this.connector + .handle_stream_error(&e, this.last_event_id.as_deref()); + let retrying = this + .connector + .retry_connection(this.last_event_id.as_deref()); + SseAutoReconnectStreamState::Retrying { retry_times: 0, + retrying, } - } else { - tracing::debug!("sse stream terminated, no reconnect policy"); - return Poll::Ready(None); } - } - } - } - SseAutoReconnectStreamStateProj::Retrying { - retry_times, - retrying, - } => { - let retry_result = ready!(retrying.poll(cx)); - match retry_result { - Ok(new_stream) => SseAutoReconnectStreamState::Connected { stream: new_stream }, - Err(e) => { - tracing::debug!("retry sse stream error: {e}"); - *retry_times += 1; - if let Some(interval) = this.retry_policy.retry(*retry_times) { + None => { + if *this.reconnect_only_after_event_id && this.last_event_id.is_none() { + tracing::debug!( + "sse response ended before an event ID was received; cannot resume" + ); + this.state.set(SseAutoReconnectStreamState::Terminated); + return Poll::Ready(None); + } + // Per SEP-1699, a graceful stream close is + // reconnectable. If the server sent a `retry` field + // we MUST wait that long before reconnecting. let interval = this .server_retry_interval - .map(|server_retry_interval| server_retry_interval.max(interval)) - .unwrap_or(interval); - let sleep = tokio::time::sleep(interval); - SseAutoReconnectStreamState::WaitingNextRetry { - sleep, - retry_times: *retry_times, + .take() + .or_else(|| this.retry_policy.retry(0)); + if let Some(interval) = interval { + tracing::debug!( + ?interval, + "sse stream ended gracefully, reconnecting" + ); + SseAutoReconnectStreamState::WaitingNextRetry { + sleep: tokio::time::sleep(interval), + retry_times: 0, + } + } else { + tracing::debug!("sse stream terminated, no reconnect policy"); + return Poll::Ready(None); } - } else { - tracing::error!("sse stream error: {e}, max retry times reached"); - this.state.set(SseAutoReconnectStreamState::Terminated); - return Poll::Ready(Some(Err(e))); } } } - } - SseAutoReconnectStreamStateProj::WaitingNextRetry { sleep, retry_times } => { - ready!(sleep.poll(cx)); - let retrying = this - .connector - .retry_connection(this.last_event_id.as_deref()); - let retry_times = *retry_times; - SseAutoReconnectStreamState::Retrying { + SseAutoReconnectStreamStateProj::Retrying { retry_times, retrying, + } => { + let retry_result = ready!(retrying.poll(cx)); + match retry_result { + Ok(new_stream) => { + SseAutoReconnectStreamState::Connected { stream: new_stream } + } + Err(e) => { + tracing::debug!("retry sse stream error: {e}"); + *retry_times += 1; + if let Some(interval) = this.retry_policy.retry(*retry_times) { + let interval = this + .server_retry_interval + .map(|server_retry_interval| { + server_retry_interval.max(interval) + }) + .unwrap_or(interval); + let sleep = tokio::time::sleep(interval); + SseAutoReconnectStreamState::WaitingNextRetry { + sleep, + retry_times: *retry_times, + } + } else { + tracing::error!("sse stream error: {e}, max retry times reached"); + this.state.set(SseAutoReconnectStreamState::Terminated); + return Poll::Ready(Some(Err(e))); + } + } + } } + SseAutoReconnectStreamStateProj::WaitingNextRetry { sleep, retry_times } => { + ready!(sleep.poll(cx)); + let retrying = this + .connector + .retry_connection(this.last_event_id.as_deref()); + let retry_times = *retry_times; + SseAutoReconnectStreamState::Retrying { + retry_times, + retrying, + } + } + SseAutoReconnectStreamStateProj::Terminated => { + return Poll::Ready(None); + } + }; + // update the state + this.state.set(next_state); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + #[derive(Debug, Error)] + enum TestReconnectError { + #[error("SSE stream error: {0}")] + Sse(SseError), + #[error("unexpected reconnect")] + Reconnect, + } + + struct CountingReconnect { + attempts: Arc, + } + + impl SseStreamReconnect for CountingReconnect { + type Error = TestReconnectError; + type Future = futures::future::Ready>; + + fn retry_connection(&mut self, _last_event_id: Option<&str>) -> Self::Future { + self.attempts.fetch_add(1, Ordering::Relaxed); + futures::future::ready(Err(TestReconnectError::Reconnect)) + } + + fn map_fatal_stream_error(&mut self, error: SseError) -> Option { + Some(TestReconnectError::Sse(error)) + } + } + + #[tokio::test] + async fn bounded_sse_stream_rejects_unterminated_event_over_limit() { + let source = futures::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"data: aaaaa")), + Ok(Bytes::from_static(b"aaaaaa")), + ]); + let mut stream = bounded_sse_stream(source, 16); + + let error = stream.next().await.unwrap().unwrap_err(); + + assert!( + error.to_string().contains("maximum size of 16 bytes"), + "unexpected error: {error}" + ); + } + + #[tokio::test] + async fn bounded_sse_stream_resets_limit_after_event_terminator() { + let source = futures::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"data: a\r")), + Ok(Bytes::from_static(b"\n\r\ndata: b\n\n")), + ]); + let mut stream = bounded_sse_stream(source, 8); + + let first = stream.next().await.unwrap().unwrap(); + let second = stream.next().await.unwrap().unwrap(); + + assert_eq!( + (first.data.as_deref(), second.data.as_deref()), + (Some("a"), Some("b")) + ); + } + + #[tokio::test] + async fn bounded_sse_stream_passes_event_at_exact_limit() { + let source = + futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"data: ab\n\n"))]); + let mut stream = bounded_sse_stream(source, 9); + + let event = stream.next().await.unwrap().unwrap(); + assert_eq!(event.data.as_deref(), Some("ab")); + } + + #[tokio::test] + async fn bounded_sse_stream_rejects_oversize_split_across_many_chunks() { + let source = futures::stream::iter( + std::iter::repeat_with(|| Ok::<_, std::io::Error>(Bytes::from_static(b"data: x"))) + .take(20), + ); + let mut stream = bounded_sse_stream(source, 32); + + let mut found_error = false; + while let Some(item) = stream.next().await { + if item.is_err() { + found_error = true; + break; } - SseAutoReconnectStreamStateProj::Terminated => { - return Poll::Ready(None); - } + } + assert!(found_error, "expected oversize error"); + } + + #[tokio::test] + async fn bounded_sse_stream_handles_crlf_split_across_chunks() { + let source = futures::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"data: hello\r")), + Ok(Bytes::from_static(b"\n\ndata: world\n\n")), + ]); + let mut stream = bounded_sse_stream(source, 64); + + let first = stream.next().await.unwrap().unwrap(); + let second = stream.next().await.unwrap().unwrap(); + + assert_eq!(first.data.as_deref(), Some("hello")); + assert_eq!(second.data.as_deref(), Some("world")); + } + + #[tokio::test] + async fn bounded_sse_stream_discards_completed_comment_lines_from_limit() { + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static( + b": ping\n: ping\n: ping\n", + ))]); + let mut stream = bounded_sse_stream(source, 6); + + assert!(stream.next().await.is_none()); + } + + #[tokio::test] + async fn bounded_sse_stream_comments_do_not_reset_accumulated_data() { + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static( + b"data: a\n: ping\ndata: b\n", + ))]); + let mut stream = bounded_sse_stream(source, 14); + + assert!(stream.next().await.unwrap().is_err()); + } + + #[tokio::test] + async fn bounded_sse_stream_counts_multiline_data_join_newlines() { + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static( + b"data: aaa\ndata: bbb\n\n", + ))]); + let mut stream = bounded_sse_stream(source, 18); + + let error = stream.next().await.unwrap().unwrap_err(); + assert!( + error.to_string().contains("maximum size"), + "unexpected error: {error}" + ); + } + + #[tokio::test] + async fn bounded_sse_stream_propagates_source_error() { + let source = futures::stream::iter([Err::(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "connection reset", + ))]); + let mut stream = bounded_sse_stream(source, 1024); + + let error = stream.next().await.unwrap().unwrap_err(); + assert!( + error.to_string().contains("connection reset"), + "unexpected error: {error}" + ); + } + + #[tokio::test] + async fn is_event_too_large_error_detects_oversize() { + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(vec![b'A'; 100]))]); + let mut stream = bounded_sse_stream(source, 8); + + let error = stream.next().await.unwrap().unwrap_err(); + assert!(is_event_too_large_error(&error)); + } + + #[tokio::test] + async fn is_event_too_large_error_rejects_other_errors() { + let source = + futures::stream::iter([Err::(std::io::Error::other("something else"))]); + let mut stream = bounded_sse_stream(source, 1024); + + let error = stream.next().await.unwrap().unwrap_err(); + assert!(!is_event_too_large_error(&error)); + } + + #[tokio::test] + async fn oversized_event_returns_error_without_reconnecting() { + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(vec![b'A'; 100]))]); + let attempts = Arc::new(AtomicUsize::new(0)); + let connector = CountingReconnect { + attempts: attempts.clone(), + }; + let stream = SseAutoReconnectStream::new( + bounded_sse_stream(source, 8), + connector, + Arc::new(NeverRetry), + ); + let mut stream = std::pin::pin!(stream); + + let result = stream.next().await; + + assert!( + matches!(result, Some(Err(TestReconnectError::Sse(_)))) + && attempts.load(Ordering::Relaxed) == 0 + ); + } + + #[tokio::test] + async fn skipped_events_do_not_grow_the_stack() { + // Every frame here fails to deserialize, so each one takes the "skip this + // event" path. The inner stream is always ready, so nothing yields in + // between: when that path recursed into poll_next instead of looping, this + // overflowed the stack and aborted the process. + const SKIPPED_EVENTS: usize = 50_000; + let mut payload = Vec::with_capacity(SKIPPED_EVENTS * 9); + for _ in 0..SKIPPED_EVENTS { + payload.extend_from_slice(b"data: x\n\n"); + } + + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(payload))]); + let attempts = Arc::new(AtomicUsize::new(0)); + let connector = CountingReconnect { + attempts: attempts.clone(), + }; + let stream = SseAutoReconnectStream::new( + bounded_sse_stream(source, 1024), + connector, + Arc::new(NeverRetry), + ); + let mut stream = std::pin::pin!(stream); + + assert!(stream.next().await.is_none()); + assert_eq!(attempts.load(Ordering::Relaxed), 0); + } + + #[tokio::test] + async fn response_without_event_id_does_not_reconnect() { + let attempts = Arc::new(AtomicUsize::new(0)); + let connector = CountingReconnect { + attempts: attempts.clone(), + }; + let stream = SseAutoReconnectStream::new_after_event_id( + futures::stream::empty().boxed(), + connector, + Arc::new(FixedInterval { + max_times: Some(1), + duration: Duration::ZERO, + }), + ); + let mut stream = std::pin::pin!(stream); + + assert!(stream.next().await.is_none()); + assert_eq!(attempts.load(Ordering::Relaxed), 0); + } + + #[test] + fn exponential_backoff_saturates_at_high_retry_counts() { + // With `max_times` unset, `current_times` can reach the bit width. The old + // `2u32.pow(current_times)` panicked in debug builds and wrapped in release; + // the saturating implementation must return a monotonic, non-zero delay instead. + let policy = ExponentialBackoff { + max_times: None, + base_duration: Duration::from_millis(1), + max_delay: None, + }; + let mut previous = Duration::ZERO; + for current_times in [31usize, 32, 63, 64, 100] { + let delay = policy + .retry(current_times) + .expect("unbounded policy never gives up"); + assert!( + !delay.is_zero(), + "delay must stay non-zero at {current_times}" + ); + assert!( + delay >= previous, + "delay must stay monotonic at {current_times}" + ); + previous = delay; + } + } + + #[test] + fn exponential_backoff_caps_delay_at_max_delay() { + // An explicit cap keeps the unbounded doubling policy from producing decades-long + // sleeps once the multiplier saturates. The delay must grow monotonically, stop at + // the configured ceiling, and never exceed it. + let policy = ExponentialBackoff { + max_times: None, + base_duration: Duration::from_secs(1), + max_delay: Some(Duration::from_secs(30)), + }; + let mut previous = Duration::ZERO; + for current_times in [0usize, 1, 2, 3, 4, 5, 10, 32, 64, 100] { + let delay = policy + .retry(current_times) + .expect("unbounded policy never gives up"); + assert!( + delay >= previous, + "delay must stay monotonic at {current_times}" + ); + assert!( + delay <= Duration::from_secs(30), + "delay must respect max_delay at {current_times}" + ); + previous = delay; + } + // Beyond the ceiling the delay stays pinned at max_delay. + assert_eq!( + policy.retry(100).expect("never gives up"), + Duration::from_secs(30) + ); + } + + #[test] + fn exponential_backoff_respects_max_times() { + let policy = ExponentialBackoff { + max_times: Some(3), + base_duration: Duration::from_millis(1), + max_delay: None, }; - // update the state - this.state.set(next_state); - self.poll_next(cx) + assert!(policy.retry(0).is_some()); + assert!(policy.retry(2).is_some()); + assert!(policy.retry(3).is_none()); } } diff --git a/crates/rmcp/src/transport/common/http_header.rs b/crates/rmcp/src/transport/common/http_header.rs index 283f0da..6c2abbc 100644 --- a/crates/rmcp/src/transport/common/http_header.rs +++ b/crates/rmcp/src/transport/common/http_header.rs @@ -4,6 +4,15 @@ pub const HEADER_MCP_PROTOCOL_VERSION: &str = "MCP-Protocol-Version"; pub const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream"; pub const JSON_MIME_TYPE: &str = "application/json"; +// SEP-2243 standard headers, gated on protocol version >= 2026-07-28. +pub const HEADER_MCP_METHOD: &str = "Mcp-Method"; +pub const HEADER_MCP_NAME: &str = "Mcp-Name"; +pub const HEADER_MCP_PARAM_PREFIX: &str = "Mcp-Param-"; + +/// Sentinel wrapping a Base64-encoded SEP-2243 header value (`=?base64??=`). +pub const BASE64_HEADER_PREFIX: &str = "=?base64?"; +pub const BASE64_HEADER_SUFFIX: &str = "?="; + /// Reserved headers that must not be overridden by user-supplied custom headers. /// `MCP-Protocol-Version` is in this list but is allowed through because the worker /// injects it after initialization. diff --git a/crates/rmcp/src/transport/common/mcp_headers.rs b/crates/rmcp/src/transport/common/mcp_headers.rs new file mode 100644 index 0000000..cc55068 --- /dev/null +++ b/crates/rmcp/src/transport/common/mcp_headers.rs @@ -0,0 +1,658 @@ +//! SEP-2243 HTTP header standardization. +//! +//! Builds and validates the `Mcp-Method`, `Mcp-Name`, and `Mcp-Param-*` headers +//! so middle boxes can route Streamable HTTP traffic without parsing the body. +//! All emission/validation is gated by the negotiated protocol version +//! (`>= ProtocolVersion::STANDARD_HEADERS`) at the call sites. + +// Which helpers are reachable depends on the client/server feature combination, +// mirroring `server_side_http`. +#![allow(dead_code)] + +use serde_json::Value; + +use super::http_header::{ + BASE64_HEADER_PREFIX, BASE64_HEADER_SUFFIX, HEADER_MCP_METHOD, HEADER_MCP_NAME, + HEADER_MCP_PARAM_PREFIX, +}; +use crate::model::JsonObject; + +/// Methods whose `Mcp-Name` is sourced from `params.name`. +const NAME_FROM_NAME: &[&str] = &["tools/call", "prompts/get"]; +/// Methods whose `Mcp-Name` is sourced from `params.uri`. +const NAME_FROM_URI: &[&str] = &[ + "resources/read", + "resources/subscribe", + "resources/unsubscribe", +]; +/// Methods whose `Mcp-Name` is sourced from `params.taskId` (SEP-2663 Tasks +/// extension): allows intermediaries to route task polling to the server +/// instance holding the task's state. +const NAME_FROM_TASK_ID: &[&str] = &["tasks/get", "tasks/update", "tasks/cancel"]; + +/// Returns the `Mcp-Name` value for a request, if the method carries one. +fn extract_name(method: &str, params: Option<&Value>) -> Option { + let params = params?; + let key = if NAME_FROM_NAME.contains(&method) { + "name" + } else if NAME_FROM_URI.contains(&method) { + "uri" + } else if NAME_FROM_TASK_ID.contains(&method) { + "taskId" + } else { + return None; + }; + params.get(key)?.as_str().map(str::to_owned) +} + +/// Converts a JSON primitive to its SEP-2243 string form. Non-primitives yield `None`. +fn primitive_to_string(value: &Value) -> Option { + match value { + Value::String(s) => Some(s.clone()), + Value::Bool(b) => Some(b.to_string()), + Value::Number(n) => Some(n.to_string()), + _ => None, + } +} + +/// True if `value` must be Base64-wrapped to survive as an HTTP header value: +/// leading/trailing space or tab, control/non-ASCII characters, or a value that +/// already looks like the `=?base64?...?=` sentinel. +#[cfg(feature = "client-side-sse")] +fn requires_base64(value: &str) -> bool { + if value.is_empty() { + return false; + } + let bytes = value.as_bytes(); + if matches!(bytes.first(), Some(b' ' | b'\t')) || matches!(bytes.last(), Some(b' ' | b'\t')) { + return true; + } + if value + .chars() + .any(|c| (c as u32) < 0x20 || (c as u32) > 0x7E) + { + return true; + } + value.starts_with(BASE64_HEADER_PREFIX) && value.ends_with(BASE64_HEADER_SUFFIX) +} + +/// RFC 9110 §5.6.2 token character. +#[cfg(feature = "client-side-sse")] +fn is_tchar(c: char) -> bool { + c.is_ascii_alphanumeric() + || matches!( + c, + '!' | '#' + | '$' + | '%' + | '&' + | '\'' + | '*' + | '+' + | '-' + | '.' + | '^' + | '_' + | '`' + | '|' + | '~' + ) +} + +/// Top-level properties carrying an `x-mcp-header` annotation, as `(property, header)` pairs. +fn param_header_annotations(input_schema: &JsonObject) -> Vec<(String, String)> { + let mut out = Vec::new(); + if let Some(Value::Object(props)) = input_schema.get("properties") { + for (prop, schema) in props { + if let Some(Value::String(header)) = schema.get("x-mcp-header") + && !header.is_empty() + { + out.push((prop.clone(), header.clone())); + } + } + } + out +} + +/// Validates the `x-mcp-header` annotations in a tool input schema. +/// +/// Annotations must be non-empty RFC 9110 tokens, case-insensitively unique, +/// applied only to top-level primitive (`string`/`integer`/`boolean`) properties. +/// Returns the offending reason on the first violation. +#[cfg(feature = "client-side-sse")] +pub(crate) fn validate_param_header_annotations(input_schema: &JsonObject) -> Result<(), String> { + let Some(Value::Object(props)) = input_schema.get("properties") else { + return Ok(()); + }; + let mut seen = std::collections::HashSet::new(); + for (prop, schema) in props { + reject_nested_annotations(schema, prop)?; + let Some(raw) = schema.get("x-mcp-header") else { + continue; + }; + let Value::String(header) = raw else { + return Err(format!("property `{prop}`: x-mcp-header must be a string")); + }; + if header.is_empty() { + return Err(format!("property `{prop}`: x-mcp-header must not be empty")); + } + if !header.chars().all(is_tchar) { + return Err(format!( + "property `{prop}`: x-mcp-header `{header}` is not a valid HTTP token" + )); + } + if !seen.insert(header.to_ascii_lowercase()) { + return Err(format!( + "property `{prop}`: duplicate x-mcp-header `{header}` (case-insensitive)" + )); + } + match schema.get("type").and_then(Value::as_str) { + Some("string" | "integer" | "boolean") => {} + other => { + return Err(format!( + "property `{prop}`: x-mcp-header requires a primitive type \ + (string/integer/boolean), got {other:?}" + )); + } + } + } + Ok(()) +} + +/// Rejects `x-mcp-header` on nested properties (only top-level promotion is supported). +#[cfg(feature = "client-side-sse")] +fn reject_nested_annotations(schema: &Value, path: &str) -> Result<(), String> { + if let Some(Value::Object(nested)) = schema.get("properties") { + for (key, value) in nested { + if value.get("x-mcp-header").is_some() { + return Err(format!( + "property `{path}.{key}`: x-mcp-header is not supported on nested properties" + )); + } + reject_nested_annotations(value, &format!("{path}.{key}"))?; + } + } + Ok(()) +} + +/// Wraps a value as `=?base64??=` when it cannot travel as a bare header value. +#[cfg(feature = "client-side-sse")] +fn encode_header_value(value: &str) -> String { + use base64::{Engine, prelude::BASE64_STANDARD}; + if requires_base64(value) { + format!( + "{BASE64_HEADER_PREFIX}{}{BASE64_HEADER_SUFFIX}", + BASE64_STANDARD.encode(value) + ) + } else { + value.to_owned() + } +} + +/// Reverses [`encode_header_value`]. Returns `None` if the sentinel wraps invalid Base64/UTF-8. +#[cfg(feature = "server-side-http")] +fn decode_header_value(value: &str) -> Option { + use base64::{Engine, prelude::BASE64_STANDARD}; + match value + .strip_prefix(BASE64_HEADER_PREFIX) + .and_then(|inner| inner.strip_suffix(BASE64_HEADER_SUFFIX)) + { + Some(inner) => { + let bytes = BASE64_STANDARD.decode(inner).ok()?; + String::from_utf8(bytes).ok() + } + None => Some(value.to_owned()), + } +} + +/// Builds the SEP-2243 headers for an outgoing request from its JSON form. +/// +/// `tool_schema` is the cached input schema of the called tool, used to promote +/// annotated `tools/call` arguments to `Mcp-Param-*` headers. +#[cfg(feature = "client-side-sse")] +pub(crate) fn standard_request_headers( + request: &Value, + tool_schema: Option<&JsonObject>, +) -> Vec<(http::HeaderName, http::HeaderValue)> { + use http::{HeaderName, HeaderValue}; + + let mut out = Vec::new(); + let Some(method) = request.get("method").and_then(Value::as_str) else { + return out; + }; + let params = request.get("params"); + + let mut push = |name: &str, value: &str| { + if let (Ok(name), Ok(value)) = ( + HeaderName::from_bytes(name.as_bytes()), + HeaderValue::from_str(value), + ) { + out.push((name, value)); + } + }; + + push(HEADER_MCP_METHOD, method); + if let Some(name) = extract_name(method, params) { + push(HEADER_MCP_NAME, &encode_header_value(&name)); + } + + if method == "tools/call" + && let (Some(schema), Some(arguments)) = + (tool_schema, params.and_then(|p| p.get("arguments"))) + { + for (prop, header) in param_header_annotations(schema) { + let Some(arg) = arguments.get(&prop) else { + continue; + }; + let Some(encoded) = primitive_to_string(arg).map(|s| encode_header_value(&s)) else { + continue; + }; + push(&format!("{HEADER_MCP_PARAM_PREFIX}{header}"), &encoded); + } + } + out +} + +/// Validates incoming SEP-2243 headers against the request body. +/// +/// Returns `Err(reason)` when a required header is missing or its value does not +/// match the body; the caller maps this to a JSON-RPC `-32020` error (HTTP 400). +#[cfg(feature = "server-side-http")] +pub(crate) fn validate_request_headers( + headers: &http::HeaderMap, + request: &Value, + tool_schema: Option<&JsonObject>, +) -> Result<(), String> { + let Some(method) = request.get("method").and_then(Value::as_str) else { + return Ok(()); + }; + let params = request.get("params"); + + let header_method = header_str(headers, HEADER_MCP_METHOD); + match header_method { + None => return Err("missing required Mcp-Method header".to_owned()), + Some(value) if value != method => { + return Err(format!( + "Mcp-Method header `{value}` does not match body method `{method}`" + )); + } + Some(_) => {} + } + + if let Some(expected) = extract_name(method, params) { + match header_str(headers, HEADER_MCP_NAME) { + None => return Err(format!("missing required Mcp-Name header for `{method}`")), + Some(raw) => { + let decoded = decode_header_value(raw) + .ok_or_else(|| "Mcp-Name header is not valid Base64".to_owned())?; + if decoded != expected { + return Err(format!( + "Mcp-Name header `{decoded}` does not match body value `{expected}`" + )); + } + } + } + } + + if method == "tools/call" + && let Some(schema) = tool_schema + { + let arguments = params.and_then(|p| p.get("arguments")); + for (prop, header) in param_header_annotations(schema) { + let full = format!("{HEADER_MCP_PARAM_PREFIX}{header}"); + let header_value = header_str(headers, &full); + let arg = arguments.and_then(|a| a.get(&prop)); + let body_value = arg.filter(|v| !v.is_null()).and_then(primitive_to_string); + + match (header_value, body_value) { + (None, None) => {} + (Some(_), None) => { + return Err(format!( + "unexpected {full} header for absent or null `{prop}`" + )); + } + (None, Some(_)) => { + return Err(format!("missing {full} header for `{prop}`")); + } + (Some(raw), Some(expected)) => { + let decoded = decode_header_value(raw) + .ok_or_else(|| format!("{full} header is not valid Base64"))?; + if decoded != expected { + return Err(format!( + "{full} header `{decoded}` does not match body value `{expected}`" + )); + } + } + } + } + } + Ok(()) +} + +/// Case-insensitive header lookup returning the value as `&str`, if present and valid UTF-8. +#[cfg(feature = "server-side-http")] +fn header_str<'a>(headers: &'a http::HeaderMap, name: &str) -> Option<&'a str> { + headers.get(name).and_then(|value| value.to_str().ok()) +} + +#[cfg(all(test, feature = "client-side-sse", feature = "server-side-http"))] +mod tests { + use std::collections::HashMap; + + use http::{HeaderMap, HeaderName, HeaderValue}; + use serde_json::json; + + use super::*; + + fn schema_with(properties: serde_json::Value) -> JsonObject { + json!({ "type": "object", "properties": properties }) + .as_object() + .unwrap() + .clone() + } + + fn header_map(pairs: &[(&str, &str)]) -> HeaderMap { + let mut map = HeaderMap::new(); + for (name, value) in pairs { + map.insert( + HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_str(value).unwrap(), + ); + } + map + } + + fn assert_wrapped(value: &str) { + let encoded = encode_header_value(value); + assert!( + encoded.starts_with(BASE64_HEADER_PREFIX), + "expected {value:?} to be Base64-wrapped, got {encoded:?}" + ); + } + + mod encode_header_value { + use super::*; + + #[test] + fn passes_plain_ascii_through() { + assert_eq!(encode_header_value("us-west1"), "us-west1"); + } + + #[test] + fn passes_internal_spaces_through() { + assert_eq!(encode_header_value("a b c"), "a b c"); + } + + #[test] + fn wraps_non_ascii() { + assert_wrapped("café"); + } + + #[test] + fn wraps_leading_whitespace() { + assert_wrapped(" padded"); + } + + #[test] + fn wraps_trailing_whitespace() { + assert_wrapped("trailing "); + } + + #[test] + fn wraps_control_characters() { + assert_wrapped("line1\nline2"); + } + + #[test] + fn wraps_crlf_injection_attempt() { + assert_wrapped("a\r\nEvil: 1"); + } + + #[test] + fn wraps_sentinel_collision() { + assert_wrapped(&format!("{BASE64_HEADER_PREFIX}x{BASE64_HEADER_SUFFIX}")); + } + } + + mod decode_header_value { + use super::*; + + #[test] + fn round_trips_with_encode() { + for value in ["us-west1", "café", " padded ", "line1\nline2", "true", "42"] { + let encoded = encode_header_value(value); + assert_eq!( + decode_header_value(&encoded).as_deref(), + Some(value), + "round-trip failed for {value:?}" + ); + } + } + + #[test] + fn rejects_invalid_base64() { + let bad = format!("{BASE64_HEADER_PREFIX}!!!not-base64!!!{BASE64_HEADER_SUFFIX}"); + assert_eq!(decode_header_value(&bad), None); + } + } + + mod extract_name { + use super::*; + + #[test] + fn from_name_for_tools_call() { + let params = json!({ "name": "my_tool" }); + assert_eq!( + extract_name("tools/call", Some(¶ms)).as_deref(), + Some("my_tool") + ); + } + + #[test] + fn from_name_for_prompts_get() { + let params = json!({ "name": "my_prompt" }); + assert_eq!( + extract_name("prompts/get", Some(¶ms)).as_deref(), + Some("my_prompt") + ); + } + + #[test] + fn from_uri_for_resources_read() { + let params = json!({ "uri": "file:///x" }); + assert_eq!( + extract_name("resources/read", Some(¶ms)).as_deref(), + Some("file:///x") + ); + } + + #[test] + fn none_for_unrelated_method() { + let params = json!({ "name": "my_tool" }); + assert_eq!(extract_name("ping", Some(¶ms)), None); + } + + #[test] + fn none_when_params_absent() { + assert_eq!(extract_name("tools/call", None), None); + } + } + + mod validate_param_header_annotations { + use super::*; + + #[test] + fn accepts_primitive_types() { + let schema = schema_with(json!({ + "region": { "type": "string", "x-mcp-header": "Region" }, + "count": { "type": "integer", "x-mcp-header": "Count" }, + "flag": { "type": "boolean", "x-mcp-header": "Flag" }, + })); + assert!(validate_param_header_annotations(&schema).is_ok()); + } + + #[test] + fn rejects_number_type() { + let schema = schema_with(json!({ "n": { "type": "number", "x-mcp-header": "N" } })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_complex_type() { + let schema = schema_with(json!({ "a": { "type": "array", "x-mcp-header": "A" } })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_empty_header_name() { + let schema = schema_with(json!({ "r": { "type": "string", "x-mcp-header": "" } })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_non_token_header_name() { + let schema = + schema_with(json!({ "r": { "type": "string", "x-mcp-header": "bad:name" } })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_case_insensitive_duplicate() { + let schema = schema_with(json!({ + "a": { "type": "string", "x-mcp-header": "Region" }, + "b": { "type": "string", "x-mcp-header": "region" }, + })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_nested_annotation() { + let schema = schema_with(json!({ + "outer": { + "type": "object", + "properties": { "inner": { "type": "string", "x-mcp-header": "Inner" } } + } + })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + } + + mod standard_request_headers { + use super::*; + + fn tools_call_headers() -> HashMap { + let schema = schema_with(json!({ + "region": { "type": "string", "x-mcp-header": "Region" }, + })); + let request = json!({ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": { "name": "deploy", "arguments": { "region": "us-west1" } } + }); + super::super::standard_request_headers(&request, Some(&schema)) + .into_iter() + .map(|(name, value)| (name.as_str().to_owned(), value.to_str().unwrap().to_owned())) + .collect() + } + + #[test] + fn sets_method_header() { + assert_eq!( + tools_call_headers().get("mcp-method").map(String::as_str), + Some("tools/call") + ); + } + + #[test] + fn sets_name_header() { + assert_eq!( + tools_call_headers().get("mcp-name").map(String::as_str), + Some("deploy") + ); + } + + #[test] + fn sets_annotated_param_header() { + assert_eq!( + tools_call_headers() + .get("mcp-param-region") + .map(String::as_str), + Some("us-west1") + ); + } + } + + mod validate_request_headers { + use super::*; + + fn tools_call_request() -> Value { + json!({ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": { "name": "deploy" } + }) + } + + #[test] + fn accepts_matching_method_and_name() { + let headers = header_map(&[("Mcp-Method", "tools/call"), ("Mcp-Name", "deploy")]); + assert!(validate_request_headers(&headers, &tools_call_request(), None).is_ok()); + } + + #[test] + fn rejects_method_mismatch() { + let headers = header_map(&[("Mcp-Method", "tools/list"), ("Mcp-Name", "deploy")]); + assert!(validate_request_headers(&headers, &tools_call_request(), None).is_err()); + } + + #[test] + fn rejects_missing_method() { + let headers = header_map(&[("Mcp-Name", "deploy")]); + assert!(validate_request_headers(&headers, &tools_call_request(), None).is_err()); + } + + #[test] + fn rejects_name_mismatch() { + let headers = header_map(&[("Mcp-Method", "tools/call"), ("Mcp-Name", "other")]); + assert!(validate_request_headers(&headers, &tools_call_request(), None).is_err()); + } + + #[test] + fn rejects_missing_name() { + let headers = header_map(&[("Mcp-Method", "tools/call")]); + assert!(validate_request_headers(&headers, &tools_call_request(), None).is_err()); + } + + #[test] + fn accepts_matching_param() { + let schema = schema_with(json!({ + "region": { "type": "string", "x-mcp-header": "Region" }, + })); + let request = json!({ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": { "name": "deploy", "arguments": { "region": "us-west1" } } + }); + let headers = header_map(&[ + ("Mcp-Method", "tools/call"), + ("Mcp-Name", "deploy"), + ("Mcp-Param-Region", "us-west1"), + ]); + assert!(validate_request_headers(&headers, &request, Some(&schema)).is_ok()); + } + + #[test] + fn rejects_param_mismatch() { + let schema = schema_with(json!({ + "region": { "type": "string", "x-mcp-header": "Region" }, + })); + let request = json!({ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": { "name": "deploy", "arguments": { "region": "us-west1" } } + }); + let headers = header_map(&[ + ("Mcp-Method", "tools/call"), + ("Mcp-Name", "deploy"), + ("Mcp-Param-Region", "eu-central1"), + ]); + assert!(validate_request_headers(&headers, &request, Some(&schema)).is_err()); + } + } +} diff --git a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs index 37c2b08..1b47560 100644 --- a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs @@ -1,16 +1,19 @@ use std::{borrow::Cow, collections::HashMap, sync::Arc}; -use futures::{StreamExt, stream::BoxStream}; +use futures::stream::BoxStream; use http::{HeaderName, HeaderValue, header::WWW_AUTHENTICATE}; use reqwest::header::ACCEPT; -use sse_stream::{Sse, SseStream}; +use sse_stream::Sse; use crate::{ model::{ClientJsonRpcMessage, JsonRpcMessage, ServerJsonRpcMessage}, transport::{ - common::http_header::{ - EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE, - extract_scope_from_header, validate_custom_header, + common::{ + client_side_sse::{DEFAULT_MAX_SSE_EVENT_SIZE, bounded_sse_stream}, + http_header::{ + EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE, + extract_scope_from_header, validate_custom_header, + }, }, streamable_http_client::*, }, @@ -49,15 +52,37 @@ impl StreamableHttpClient for reqwest::Client { async fn get_stream( &self, uri: Arc, - session_id: Arc, + session_id: Option>, last_event_id: Option, auth_token: Option, custom_headers: HashMap, + ) -> Result>, StreamableHttpError> { + self.get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + auth_token, + custom_headers, + DEFAULT_MAX_SSE_EVENT_SIZE, + ) + .await + } + + async fn get_stream_with_max_sse_event_size( + &self, + uri: Arc, + session_id: Option>, + last_event_id: Option, + auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, ) -> Result>, StreamableHttpError> { let mut request_builder = self .get(uri.as_ref()) - .header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", ")) - .header(HEADER_SESSION_ID, session_id.as_ref()); + .header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", ")); + if let Some(session_id) = session_id { + request_builder = request_builder.header(HEADER_SESSION_ID, session_id.as_ref()); + } if let Some(last_event_id) = last_event_id { request_builder = request_builder.header(HEADER_LAST_EVENT_ID, last_event_id); } @@ -69,6 +94,37 @@ impl StreamableHttpClient for reqwest::Client { if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED { return Err(StreamableHttpError::ServerDoesNotSupportSse); } + if response.status() == reqwest::StatusCode::UNAUTHORIZED + && let Some(header) = response.headers().get(WWW_AUTHENTICATE) + { + let header = header + .to_str() + .map_err(|_| { + StreamableHttpError::UnexpectedServerResponse(Cow::from( + "invalid www-authenticate header value", + )) + })? + .to_string(); + return Err(StreamableHttpError::AuthRequired(AuthRequiredError { + www_authenticate_header: header, + })); + } + if response.status() == reqwest::StatusCode::FORBIDDEN + && let Some(header) = response.headers().get(WWW_AUTHENTICATE) + { + let header_str = header.to_str().map_err(|_| { + StreamableHttpError::UnexpectedServerResponse(Cow::from( + "invalid www-authenticate header value", + )) + })?; + let scope = extract_scope_from_header(header_str); + return Err(StreamableHttpError::InsufficientScope( + InsufficientScopeError { + www_authenticate_header: header_str.to_string(), + required_scope: scope, + }, + )); + } let response = response.error_for_status()?; match response.headers().get(reqwest::header::CONTENT_TYPE) { Some(ct) => { @@ -84,7 +140,7 @@ impl StreamableHttpClient for reqwest::Client { return Err(StreamableHttpError::UnexpectedContentType(None)); } } - let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed(); + let event_stream = bounded_sse_stream(response.bytes_stream(), max_sse_event_size); Ok(event_stream) } @@ -119,6 +175,26 @@ impl StreamableHttpClient for reqwest::Client { session_id: Option>, auth_token: Option, custom_headers: HashMap, + ) -> Result> { + self.post_message_with_max_sse_event_size( + uri, + message, + session_id, + auth_token, + custom_headers, + DEFAULT_MAX_SSE_EVENT_SIZE, + ) + .await + } + + async fn post_message_with_max_sse_event_size( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, ) -> Result> { let mut request = self .post(uri.as_ref()) @@ -133,36 +209,36 @@ impl StreamableHttpClient for reqwest::Client { request = request.header(HEADER_SESSION_ID, session_id.as_ref()); } let response = request.json(&message).send().await?; - if response.status() == reqwest::StatusCode::UNAUTHORIZED { - if let Some(header) = response.headers().get(WWW_AUTHENTICATE) { - let header = header - .to_str() - .map_err(|_| { - StreamableHttpError::UnexpectedServerResponse(Cow::from( - "invalid www-authenticate header value", - )) - })? - .to_string(); - return Err(StreamableHttpError::AuthRequired(AuthRequiredError { - www_authenticate_header: header, - })); - } - } - if response.status() == reqwest::StatusCode::FORBIDDEN { - if let Some(header) = response.headers().get(WWW_AUTHENTICATE) { - let header_str = header.to_str().map_err(|_| { + if response.status() == reqwest::StatusCode::UNAUTHORIZED + && let Some(header) = response.headers().get(WWW_AUTHENTICATE) + { + let header = header + .to_str() + .map_err(|_| { StreamableHttpError::UnexpectedServerResponse(Cow::from( "invalid www-authenticate header value", )) - })?; - let scope = extract_scope_from_header(header_str); - return Err(StreamableHttpError::InsufficientScope( - InsufficientScopeError { - www_authenticate_header: header_str.to_string(), - required_scope: scope, - }, - )); - } + })? + .to_string(); + return Err(StreamableHttpError::AuthRequired(AuthRequiredError { + www_authenticate_header: header, + })); + } + if response.status() == reqwest::StatusCode::FORBIDDEN + && let Some(header) = response.headers().get(WWW_AUTHENTICATE) + { + let header_str = header.to_str().map_err(|_| { + StreamableHttpError::UnexpectedServerResponse(Cow::from( + "invalid www-authenticate header value", + )) + })?; + let scope = extract_scope_from_header(header_str); + return Err(StreamableHttpError::InsufficientScope( + InsufficientScopeError { + www_authenticate_header: header_str.to_string(), + required_scope: scope, + }, + )); } let status = response.status(); if matches!( @@ -217,27 +293,44 @@ impl StreamableHttpClient for reqwest::Client { ), } } + if let Some(response) = + legacy_discover_response(&message, session_was_attached, status, &body) + { + return Ok(response); + } return Err(StreamableHttpError::UnexpectedServerResponse(Cow::Owned( format!("HTTP {status}: {body}"), ))); } match content_type.as_deref() { Some(ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => { - let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed(); + let event_stream = bounded_sse_stream(response.bytes_stream(), max_sse_event_size); Ok(StreamableHttpPostResponse::Sse(event_stream, session_id)) } Some(ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => { - // Try to parse as a valid JSON-RPC message. If the body is - // malformed (e.g. a 200 response to a notification that lacks - // an `id` field), treat it as accepted rather than failing. - match response.json::().await { - Ok(message) => Ok(StreamableHttpPostResponse::Json(message, session_id)), - Err(e) => { + // A notification/response/error POST does not await a JSON-RPC + // reply. Treat an unusable JSON body as Accepted. A request + // still needs a reply; the same body is an error so the worker + // does not wait on SSE forever. + let body = response.bytes().await?; + match serde_json::from_slice::(&body) { + Ok(parsed) => Ok(StreamableHttpPostResponse::Json(parsed, session_id)), + Err(e) + if matches!( + message, + ClientJsonRpcMessage::Notification(_) + | ClientJsonRpcMessage::Response(_) + | ClientJsonRpcMessage::Error(_) + ) => + { tracing::warn!( "could not parse JSON response as ServerJsonRpcMessage, treating as accepted: {e}" ); Ok(StreamableHttpPostResponse::Accepted) } + Err(e) => Err(StreamableHttpError::UnexpectedServerResponse(Cow::Owned( + json_rpc_parse_error_message(&e, &body), + ))), } } _ => { @@ -364,6 +457,57 @@ mod tests { assert!(parse_json_rpc_error(body).is_none()); } + #[tokio::test] + async fn post_sse_response_honors_configured_event_limit() -> anyhow::Result<()> { + use std::{collections::HashMap, net::SocketAddr, sync::Arc}; + + use axum::{Router, routing::post}; + use futures::StreamExt; + + use crate::transport::streamable_http_client::{ + StreamableHttpClient, StreamableHttpPostResponse, + }; + + let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0))).await?; + let addr = listener.local_addr()?; + let server = tokio::spawn(async move { + let app = Router::new().route( + "/mcp", + post(|| async { + ( + [(http::header::CONTENT_TYPE, "text/event-stream")], + "data: this event is too large\n", + ) + }), + ); + axum::serve(listener, app).await + }); + let message = ClientJsonRpcMessage::request( + ClientRequest::PingRequest(PingRequest::default()), + RequestId::Number(1), + ); + let client = reqwest::Client::new(); + + let response = client + .post_message_with_max_sse_event_size( + Arc::::from(format!("http://{addr}/mcp")), + message, + None, + None, + HashMap::new(), + 16, + ) + .await?; + let StreamableHttpPostResponse::Sse(mut stream, _) = response else { + anyhow::bail!("expected SSE response"); + }; + let error = stream.next().await.unwrap().unwrap_err(); + + server.abort(); + assert!(error.to_string().contains("maximum size of 16 bytes")); + Ok(()) + } + #[tokio::test] async fn default_http_client_does_not_leak_custom_headers_to_redirect_target() -> anyhow::Result<()> { diff --git a/crates/rmcp/src/transport/common/server_side_http.rs b/crates/rmcp/src/transport/common/server_side_http.rs index 39a321f..09df6d2 100644 --- a/crates/rmcp/src/transport/common/server_side_http.rs +++ b/crates/rmcp/src/transport/common/server_side_http.rs @@ -6,7 +6,7 @@ use http::Response; use http_body::Body; use http_body_util::{BodyExt, Empty, Full, combinators::BoxBody}; use sse_stream::{KeepAlive, Sse, SseBody}; -use tokio_util::sync::CancellationToken; +use tokio_util::sync::{CancellationToken, DropGuard}; use super::http_header::EVENT_STREAM_MIME_TYPE; use crate::model::{ClientJsonRpcMessage, ServerJsonRpcMessage}; @@ -57,6 +57,25 @@ impl sse_stream::Timer for TokioTimer { } } +pin_project_lite::pin_project! { + struct CancelOnDropStream { + #[pin] + inner: S, + _drop_guard: DropGuard, + } +} + +impl futures::Stream for CancelOnDropStream { + type Item = S::Item; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + self.project().inner.poll_next(cx) + } +} + #[derive(Debug, Clone, Default)] #[non_exhaustive] pub struct ServerSseMessage { @@ -100,6 +119,15 @@ impl ServerSseMessage { retry: Some(retry), } } + + /// Create a retry hint without changing the client's last event ID. + pub fn retry(retry: Duration) -> Self { + Self { + event_id: None, + message: None, + retry: Some(retry), + } + } } pub(crate) fn sse_stream_response( @@ -108,6 +136,7 @@ pub(crate) fn sse_stream_response( ct: CancellationToken, ) -> Response> { use futures::StreamExt; + let cancelled = ct.clone(); let stream = stream .map(|message| { let mut sse = if let Some(ref msg) = message.message { @@ -126,7 +155,11 @@ pub(crate) fn sse_stream_response( Result::::Ok(sse) }) - .take_until(async move { ct.cancelled().await }); + .take_until(async move { cancelled.cancelled().await }); + let stream = CancelOnDropStream { + inner: stream, + _drop_guard: ct.drop_guard(), + }; let stream = SseBody::new(stream); let stream = match keep_alive { @@ -140,6 +173,7 @@ pub(crate) fn sse_stream_response( .status(http::StatusCode::OK) .header(http::header::CONTENT_TYPE, EVENT_STREAM_MIME_TYPE) .header(http::header::CACHE_CONTROL, "no-cache") + .header("X-Accel-Buffering", "no") .body(stream) .expect("valid response") } @@ -170,31 +204,61 @@ pub(crate) fn unexpected_message_response(expect: &str) -> Response( body: B, + max_bytes: usize, ) -> Result>> where B: Body + Send + 'static, B::Error: Display, { - match body.collect().await { - Ok(bytes) => { - match serde_json::from_reader::<_, ClientJsonRpcMessage>(bytes.aggregate().reader()) { - Ok(message) => Ok(message), - Err(e) => { + let mut collected = bytes::BytesMut::new(); + let mut body = std::pin::pin!(body); + loop { + let frame = futures::future::poll_fn(|cx| body.as_mut().poll_frame(cx)).await; + match frame { + None => break, + Some(Ok(frame)) => { + let Ok(mut data) = frame.into_data() else { + continue; + }; + if data.remaining() > max_bytes.saturating_sub(collected.len()) { let response = Response::builder() - .status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE) + .status(http::StatusCode::PAYLOAD_TOO_LARGE) .body( - Full::new(Bytes::from(format!("fail to deserialize request body {e}"))) - .boxed(), + Full::new(Bytes::from(format!( + "Payload Too Large: request body exceeds {max_bytes} bytes" + ))) + .boxed(), ) .expect("valid response"); - Err(response) + return Err(response); + } + while data.has_remaining() { + let chunk = data.chunk(); + let chunk_len = chunk.len(); + collected.extend_from_slice(chunk); + data.advance(chunk_len); } } + Some(Err(e)) => { + let response = Response::builder() + .status(http::StatusCode::INTERNAL_SERVER_ERROR) + .body( + Full::new(Bytes::from(format!("Failed to read request body: {e}"))).boxed(), + ) + .expect("valid response"); + return Err(response); + } } + } + + match serde_json::from_slice::(&collected) { + Ok(message) => Ok(message), Err(e) => { let response = Response::builder() - .status(http::StatusCode::INTERNAL_SERVER_ERROR) - .body(Full::new(Bytes::from(format!("Failed to read request body: {e}"))).boxed()) + .status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE) + .body( + Full::new(Bytes::from(format!("fail to deserialize request body {e}"))).boxed(), + ) .expect("valid response"); Err(response) } @@ -203,9 +267,17 @@ where #[cfg(test)] mod tests { + use std::convert::Infallible; + + use futures::stream; + use http_body::Frame; + use http_body_util::StreamBody; + use super::*; use crate::model::{EmptyResult, JsonRpcResponse, JsonRpcVersion2_0, RequestId, ServerResult}; + const INITIALIZE_REQUEST: &str = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"0.1"}}}"#; + fn dummy_message() -> ServerJsonRpcMessage { ServerJsonRpcMessage::Response(JsonRpcResponse { jsonrpc: JsonRpcVersion2_0, @@ -245,4 +317,62 @@ mod tests { assert!(msg.message.is_none()); assert_eq!(msg.retry, Some(Duration::from_secs(5))); } + + #[tokio::test] + async fn expect_json_accepts_body_under_limit() { + let body = Full::new(Bytes::from_static(INITIALIZE_REQUEST.as_bytes())); + let result = expect_json(body, 4 * 1024 * 1024).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn expect_json_accepts_non_contiguous_body_data() { + let split = INITIALIZE_REQUEST.len() / 2; + let data = Bytes::copy_from_slice(&INITIALIZE_REQUEST.as_bytes()[..split]).chain( + Bytes::copy_from_slice(&INITIALIZE_REQUEST.as_bytes()[split..]), + ); + let result = expect_json(Full::new(data), 4 * 1024 * 1024).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn expect_json_accepts_multi_frame_body() { + let split = INITIALIZE_REQUEST.len() / 2; + let body = StreamBody::new(stream::iter([ + Ok::<_, Infallible>(Frame::data(Bytes::copy_from_slice( + &INITIALIZE_REQUEST.as_bytes()[..split], + ))), + Ok(Frame::data(Bytes::copy_from_slice( + &INITIALIZE_REQUEST.as_bytes()[split..], + ))), + ])); + let result = expect_json(body, 4 * 1024 * 1024).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn expect_json_rejects_oversized_body() { + let big_body = Full::new(Bytes::from(vec![b'x'; 128])); + let result = expect_json(big_body, 64).await; + let response = result.unwrap_err(); + assert_eq!(response.status(), http::StatusCode::PAYLOAD_TOO_LARGE); + } + + #[tokio::test] + async fn expect_json_rejects_oversized_multi_frame_body() { + let body = StreamBody::new(stream::iter([ + Ok::<_, Infallible>(Frame::data(Bytes::from_static(b"12345678"))), + Ok(Frame::data(Bytes::from_static(b"9"))), + ])); + let response = expect_json(body, 8).await.unwrap_err(); + assert_eq!(response.status(), http::StatusCode::PAYLOAD_TOO_LARGE); + } + + #[tokio::test] + async fn expect_json_returns_415_for_invalid_json_under_limit() { + let body = Full::new(Bytes::from("not valid json")); + let result = expect_json(body, 4 * 1024 * 1024).await; + let response = result.unwrap_err(); + assert_eq!(response.status(), http::StatusCode::UNSUPPORTED_MEDIA_TYPE); + } } diff --git a/crates/rmcp/src/transport/common/unix_socket.rs b/crates/rmcp/src/transport/common/unix_socket.rs index 8ea30f5..688c9ad 100644 --- a/crates/rmcp/src/transport/common/unix_socket.rs +++ b/crates/rmcp/src/transport/common/unix_socket.rs @@ -1,20 +1,23 @@ use std::{borrow::Cow, collections::HashMap, sync::Arc}; use bytes::Bytes; -use futures::{StreamExt, stream::BoxStream}; +use futures::stream::BoxStream; use http::{HeaderName, HeaderValue, Method, Request, StatusCode, header::WWW_AUTHENTICATE}; use http_body_util::{BodyExt, Full}; use hyper::body::Incoming; use hyper_util::rt::TokioIo; -use sse_stream::{Sse, SseStream}; +use sse_stream::Sse; use tokio::net::UnixStream; use crate::{ model::{ClientJsonRpcMessage, ServerJsonRpcMessage}, transport::{ - common::http_header::{ - EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE, - extract_scope_from_header, validate_custom_header, + common::{ + client_side_sse::{DEFAULT_MAX_SSE_EVENT_SIZE, bounded_sse_stream}, + http_header::{ + EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE, + extract_scope_from_header, validate_custom_header, + }, }, streamable_http_client::*, }, @@ -168,6 +171,26 @@ impl StreamableHttpClient for UnixSocketHttpClient { session_id: Option>, auth_token: Option, custom_headers: HashMap, + ) -> Result> { + self.post_message_with_max_sse_event_size( + uri, + message, + session_id, + auth_token, + custom_headers, + DEFAULT_MAX_SSE_EVENT_SIZE, + ) + .await + } + + async fn post_message_with_max_sse_event_size( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, ) -> Result> { let json_body = serde_json::to_string(&message) .map_err(|e| StreamableHttpError::Client(UnixSocketError::Json(e)))?; @@ -203,37 +226,37 @@ impl StreamableHttpClient for UnixSocketHttpClient { let status = response.status(); - if status == StatusCode::UNAUTHORIZED { - if let Some(header) = response.headers().get(WWW_AUTHENTICATE) { - let www_authenticate_header = header - .to_str() - .map_err(|_| { - StreamableHttpError::UnexpectedServerResponse(Cow::from( - "invalid www-authenticate header value", - )) - })? - .to_string(); - return Err(StreamableHttpError::AuthRequired(AuthRequiredError { - www_authenticate_header, - })); - } - } - - if status == StatusCode::FORBIDDEN { - if let Some(header) = response.headers().get(WWW_AUTHENTICATE) { - let header_str = header.to_str().map_err(|_| { + if status == StatusCode::UNAUTHORIZED + && let Some(header) = response.headers().get(WWW_AUTHENTICATE) + { + let www_authenticate_header = header + .to_str() + .map_err(|_| { StreamableHttpError::UnexpectedServerResponse(Cow::from( "invalid www-authenticate header value", )) - })?; - let scope = extract_scope_from_header(header_str); - return Err(StreamableHttpError::InsufficientScope( - InsufficientScopeError { - www_authenticate_header: header_str.to_string(), - required_scope: scope, - }, - )); - } + })? + .to_string(); + return Err(StreamableHttpError::AuthRequired(AuthRequiredError { + www_authenticate_header, + })); + } + + if status == StatusCode::FORBIDDEN + && let Some(header) = response.headers().get(WWW_AUTHENTICATE) + { + let header_str = header.to_str().map_err(|_| { + StreamableHttpError::UnexpectedServerResponse(Cow::from( + "invalid www-authenticate header value", + )) + })?; + let scope = extract_scope_from_header(header_str); + return Err(StreamableHttpError::InsufficientScope( + InsufficientScopeError { + www_authenticate_header: header_str.to_string(), + required_scope: scope, + }, + )); } if matches!(status, StatusCode::ACCEPTED | StatusCode::NO_CONTENT) { @@ -251,6 +274,11 @@ impl StreamableHttpClient for UnixSocketHttpClient { .await .map(|c| String::from_utf8_lossy(&c.to_bytes()).into_owned()) .unwrap_or_else(|_| "".to_owned()); + if let Some(response) = + legacy_discover_response(&message, session_was_attached, status, &body) + { + return Ok(response); + } return Err(StreamableHttpError::UnexpectedServerResponse(Cow::Owned( format!("HTTP {status}: {body}"), ))); @@ -282,7 +310,8 @@ impl StreamableHttpClient for UnixSocketHttpClient { match content_type { Some(ref ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => { - let sse_stream = SseStream::new(response.into_body()).boxed(); + let sse_stream = + bounded_sse_stream(response.into_body().into_data_stream(), max_sse_event_size); Ok(StreamableHttpPostResponse::Sse(sse_stream, session_id)) } Some(ref ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => { @@ -293,13 +322,23 @@ impl StreamableHttpClient for UnixSocketHttpClient { .map_err(|e| StreamableHttpError::Client(UnixSocketError::Hyper(e)))? .to_bytes(); match serde_json::from_slice::(&body) { - Ok(message) => Ok(StreamableHttpPostResponse::Json(message, session_id)), - Err(e) => { + Ok(parsed) => Ok(StreamableHttpPostResponse::Json(parsed, session_id)), + Err(e) + if matches!( + message, + ClientJsonRpcMessage::Notification(_) + | ClientJsonRpcMessage::Response(_) + | ClientJsonRpcMessage::Error(_) + ) => + { tracing::warn!( "could not parse JSON response as ServerJsonRpcMessage, treating as accepted: {e}" ); Ok(StreamableHttpPostResponse::Accepted) } + Err(e) => Err(StreamableHttpError::UnexpectedServerResponse(Cow::Owned( + json_rpc_parse_error_message(&e, &body), + ))), } } _ => Err(StreamableHttpError::UnexpectedContentType( @@ -352,11 +391,32 @@ impl StreamableHttpClient for UnixSocketHttpClient { async fn get_stream( &self, uri: Arc, - session_id: Arc, + session_id: Option>, last_event_id: Option, auth_token: Option, custom_headers: HashMap, ) -> Result>, StreamableHttpError> + { + self.get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + auth_token, + custom_headers, + DEFAULT_MAX_SSE_EVENT_SIZE, + ) + .await + } + + async fn get_stream_with_max_sse_event_size( + &self, + uri: Arc, + session_id: Option>, + last_event_id: Option, + auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, + ) -> Result>, StreamableHttpError> { let mut builder = Request::builder() .method(Method::GET) @@ -365,8 +425,11 @@ impl StreamableHttpClient for UnixSocketHttpClient { .header( http::header::ACCEPT, format!("{EVENT_STREAM_MIME_TYPE}, {JSON_MIME_TYPE}"), - ) - .header(HEADER_SESSION_ID, session_id.as_ref()); + ); + + if let Some(session_id) = session_id { + builder = builder.header(HEADER_SESSION_ID, session_id.as_ref()); + } if let Some(last_id) = last_event_id { builder = builder.header(HEADER_LAST_EVENT_ID, last_id); @@ -390,37 +453,37 @@ impl StreamableHttpClient for UnixSocketHttpClient { return Err(StreamableHttpError::ServerDoesNotSupportSse); } - if response.status() == StatusCode::UNAUTHORIZED { - if let Some(header) = response.headers().get(WWW_AUTHENTICATE) { - let www_authenticate_header = header - .to_str() - .map_err(|_| { - StreamableHttpError::UnexpectedServerResponse(Cow::from( - "invalid www-authenticate header value", - )) - })? - .to_string(); - return Err(StreamableHttpError::AuthRequired(AuthRequiredError { - www_authenticate_header, - })); - } - } - - if response.status() == StatusCode::FORBIDDEN { - if let Some(header) = response.headers().get(WWW_AUTHENTICATE) { - let header_str = header.to_str().map_err(|_| { + if response.status() == StatusCode::UNAUTHORIZED + && let Some(header) = response.headers().get(WWW_AUTHENTICATE) + { + let www_authenticate_header = header + .to_str() + .map_err(|_| { StreamableHttpError::UnexpectedServerResponse(Cow::from( "invalid www-authenticate header value", )) - })?; - let scope = extract_scope_from_header(header_str); - return Err(StreamableHttpError::InsufficientScope( - InsufficientScopeError { - www_authenticate_header: header_str.to_string(), - required_scope: scope, - }, - )); - } + })? + .to_string(); + return Err(StreamableHttpError::AuthRequired(AuthRequiredError { + www_authenticate_header, + })); + } + + if response.status() == StatusCode::FORBIDDEN + && let Some(header) = response.headers().get(WWW_AUTHENTICATE) + { + let header_str = header.to_str().map_err(|_| { + StreamableHttpError::UnexpectedServerResponse(Cow::from( + "invalid www-authenticate header value", + )) + })?; + let scope = extract_scope_from_header(header_str); + return Err(StreamableHttpError::InsufficientScope( + InsufficientScopeError { + www_authenticate_header: header_str.to_string(), + required_scope: scope, + }, + )); } if !response.status().is_success() { @@ -444,7 +507,10 @@ impl StreamableHttpClient for UnixSocketHttpClient { } } - Ok(SseStream::new(response.into_body()).boxed()) + Ok(bounded_sse_stream( + response.into_body().into_data_stream(), + max_sse_event_size, + )) } } diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index a2c1a7b..8af045f 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -1,29 +1,139 @@ -use std::{borrow::Cow, collections::HashMap, sync::Arc, time::Duration}; +use std::{ + borrow::Cow, + collections::{HashMap, HashSet, VecDeque}, + sync::Arc, + time::Duration, +}; -use futures::{Stream, StreamExt, future::BoxFuture, stream::BoxStream}; -use http::{HeaderName, HeaderValue}; +use futures::{ + Stream, StreamExt, + future::BoxFuture, + stream::{BoxStream, FuturesUnordered}, +}; +use http::{HeaderName, HeaderValue, StatusCode}; pub use sse_stream::Error as SseError; use sse_stream::Sse; use thiserror::Error; use tokio_util::sync::CancellationToken; use tracing::debug; -use super::common::client_side_sse::{ExponentialBackoff, SseRetryPolicy, SseStreamReconnect}; +use super::common::client_side_sse::{ + DEFAULT_MAX_SSE_EVENT_SIZE, ExponentialBackoff, SseRetryPolicy, SseStreamReconnect, +}; use crate::{ RoleClient, model::{ - ClientJsonRpcMessage, ClientNotification, InitializedNotification, ServerJsonRpcMessage, + ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, GetExtensions, GetMeta, + InitializedNotification, JsonObject, ProtocolVersion, RequestId, ServerJsonRpcMessage, ServerResult, }, + service::InboundStreamOrigin, transport::{ - common::client_side_sse::SseAutoReconnectStream, - worker::{Worker, WorkerQuitReason, WorkerSendRequest, WorkerTransport}, + common::{client_side_sse::SseAutoReconnectStream, mcp_headers}, + worker::{ + RequestCancellationRegistration, Worker, WorkerQuitReason, WorkerSendRequest, + WorkerTransport, + }, }, }; type BoxedSseStream = BoxStream<'static, Result>; +type SseTaskResult = (Option, Result<(), StreamableHttpError>); +const SESSION_CLEANUP_TIMEOUT: Duration = Duration::from_secs(5); -#[derive(Debug)] +fn build_request_headers( + base: &HashMap, + message: &ClientJsonRpcMessage, + tool_cache: &HashMap>, + version: &ProtocolVersion, +) -> HashMap { + use serde_json::Value; + + let mut headers = base.clone(); + if *version >= ProtocolVersion::STANDARD_HEADERS + && let Ok(value) = serde_json::to_value(message) + { + let schema = value + .get("method") + .and_then(Value::as_str) + .filter(|method| *method == "tools/call") + .and_then(|_| value.get("params")) + .and_then(|params| params.get("name")) + .and_then(Value::as_str) + .and_then(|name| tool_cache.get(name)) + .map(Arc::as_ref); + for (name, val) in mcp_headers::standard_request_headers(&value, schema) { + headers.insert(name, val); + } + } + headers +} + +fn request_version_headers( + base: &HashMap, + message: &ClientJsonRpcMessage, + fallback: &ProtocolVersion, + tool_cache: &HashMap>, +) -> (ProtocolVersion, HashMap) { + let version = match message { + ClientJsonRpcMessage::Request(request) => request + .request + .get_meta() + .protocol_version() + .unwrap_or_else(|| fallback.clone()), + _ => fallback.clone(), + }; + let mut headers = build_request_headers(base, message, tool_cache, &version); + if let Ok(value) = HeaderValue::from_str(version.as_str()) { + headers.insert(HeaderName::from_static("mcp-protocol-version"), value); + } + (version, headers) +} + +fn cache_tools_from_response( + cache: &mut HashMap>, + message: &mut ServerJsonRpcMessage, + protocol_version: &ProtocolVersion, +) { + if protocol_version < &ProtocolVersion::STANDARD_HEADERS { + return; + } + if let ServerJsonRpcMessage::Response(response) = message + && let ServerResult::ListToolsResult(list) = &mut response.result + { + list.tools.retain(|tool| { + let Err(reason) = + mcp_headers::validate_param_header_annotations(&tool.input_schema) + else { + cache.insert(tool.name.to_string(), tool.input_schema.clone()); + return true; + }; + tracing::warn!(tool = %tool.name, "rejecting invalid x-mcp-header annotations: {reason}"); + false + }); + } +} + +fn negotiate_version_headers( + init_response: &ServerJsonRpcMessage, + base: HashMap, +) -> (ProtocolVersion, HashMap) { + let mut version = ProtocolVersion::default(); + let mut headers = base; + if let ServerJsonRpcMessage::Response(response) = init_response + && let ServerResult::InitializeResult(init_result) = &response.result + { + version = init_result.protocol_version.clone(); + // HeaderName::from_static requires lowercase + if let Ok(hv) = HeaderValue::from_str(init_result.protocol_version.as_str()) { + headers.insert(HeaderName::from_static("mcp-protocol-version"), hv); + } + } + (version, headers) +} + +#[derive(Debug, Error)] +#[error("authorization required: {www_authenticate_header}")] #[non_exhaustive] pub struct AuthRequiredError { pub www_authenticate_header: String, @@ -38,7 +148,8 @@ impl AuthRequiredError { } } -#[derive(Debug)] +#[derive(Debug, Error)] +#[error("insufficient scope: {www_authenticate_header}")] #[non_exhaustive] pub struct InsufficientScopeError { pub www_authenticate_header: String, @@ -96,13 +207,64 @@ pub enum StreamableHttpError { #[error("Auth error: {0}")] Auth(#[from] crate::transport::auth::AuthError), #[error("Auth required")] - AuthRequired(AuthRequiredError), + AuthRequired(#[source] AuthRequiredError), #[error("Insufficient scope")] - InsufficientScope(InsufficientScopeError), + InsufficientScope(#[source] InsufficientScopeError), #[error("Header name '{0}' is reserved and conflicts with default headers")] ReservedHeaderConflict(String), #[error("Session expired (HTTP 404)")] SessionExpired, + /// Session recovery timed out. The server may have processed an interrupted POST. + #[error("Session recovery timed out; the server may have processed the POST")] + SessionRecoveryTimeout, + /// A cancellation or reply POST did not finish in time. + #[error("Control POST timed out")] + ControlRequestTimeout, +} + +/// Bytes of an unusable JSON body kept in +/// [`StreamableHttpError::UnexpectedServerResponse`]. +pub(crate) const JSON_RPC_BODY_PREVIEW_LEN: usize = 256; + +/// Short, lossy view of a body that failed to parse as JSON-RPC. +pub(crate) fn json_rpc_body_preview(body: &[u8]) -> Cow<'static, str> { + if body.is_empty() { + return Cow::Borrowed(""); + } + let lossy = String::from_utf8_lossy(body); + if lossy.len() <= JSON_RPC_BODY_PREVIEW_LEN { + return Cow::Owned(lossy.into_owned()); + } + let mut end = JSON_RPC_BODY_PREVIEW_LEN; + while end > 0 && !lossy.is_char_boundary(end) { + end -= 1; + } + Cow::Owned(lossy[..end].to_owned()) +} + +/// Request-POST parse failure, including a body preview so empty / HTML / `{}` +/// are distinguishable from the serde diagnostic alone. +pub(crate) fn json_rpc_parse_error_message(err: &serde_json::Error, body: &[u8]) -> String { + format!( + "could not parse JSON response as ServerJsonRpcMessage: {err}: {}", + json_rpc_body_preview(body) + ) +} + +impl StreamableHttpError { + /// The `WWW-Authenticate` challenge carried by this error, when the + /// server answered 401 ([`AuthRequired`](Self::AuthRequired)) or 403 + /// ([`InsufficientScope`](Self::InsufficientScope)). Feed it to + /// [`AuthorizationRequest::with_challenge`](crate::transport::auth::AuthorizationRequest::with_challenge) + /// to authorize reactively. + #[cfg(feature = "auth")] + pub fn auth_challenge(&self) -> Option<&str> { + match self { + Self::AuthRequired(error) => Some(&error.www_authenticate_header), + Self::InsufficientScope(error) => Some(&error.www_authenticate_header), + _ => None, + } + } } #[derive(Debug, Clone, Error)] @@ -112,7 +274,10 @@ pub enum StreamableHttpProtocolError { MissingSessionIdInResponse, } -#[allow(clippy::large_enum_variant)] +#[expect( + clippy::large_enum_variant, + reason = "boxing the streaming response would add an allocation to the common response path" +)] #[non_exhaustive] pub enum StreamableHttpPostResponse { Accepted, @@ -196,6 +361,53 @@ impl StreamableHttpPostResponse { } } +/// Convert a sessionless discovery rejection into a response the lifecycle +/// layer can classify as a legacy-server signal. +/// +/// Some legacy streamable-HTTP servers reject `server/discover` in middleware +/// before it reaches JSON-RPC dispatch. Their response may be an empty or +/// plain-text 4xx, so there is no JSON-RPC error for the client to forward. +/// Keep authentication failures and server errors on their original paths. +pub(super) fn legacy_discover_response( + message: &ClientJsonRpcMessage, + session_was_attached: bool, + status: StatusCode, + body: &str, +) -> Option { + if session_was_attached + || !status.is_client_error() + || matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) + { + return None; + } + + let ClientJsonRpcMessage::Request(request) = message else { + return None; + }; + if !matches!(request.request, ClientRequest::DiscoverRequest(_)) { + return None; + } + + let error = ErrorData::invalid_request( + format!("server/discover rejected with HTTP {status}: {body}"), + None, + ); + Some(StreamableHttpPostResponse::Json( + ServerJsonRpcMessage::error(error, Some(request.id.clone())), + None, + )) +} + +/// HTTP backend used by [`StreamableHttpClientTransport`]. +/// +/// Custom implementations that parse SSE responses must override +/// [`Self::post_message_with_max_sse_event_size`] and +/// [`Self::get_stream_with_max_sse_event_size`] to enforce the transport's +/// configured event-size limit. +/// +/// For legacy http, the transport keeps an open response stream alive until +/// its cancellation send finishes or is dropped. This lets a custom client +/// handle the cancellation using state owned by that stream. pub trait StreamableHttpClient: Clone + Send + 'static { type Error: std::error::Error + Send + Sync + 'static; fn post_message( @@ -208,6 +420,29 @@ pub trait StreamableHttpClient: Clone + Send + 'static { ) -> impl Future>> + Send + '_; + /// Send a message while enforcing a maximum raw SSE event size. + /// + /// `max_sse_event_size` is not a per-request option: it is the + /// transport-wide [`StreamableHttpClientTransportConfig::max_sse_event_size`] + /// value, passed identically on every call because the limit must be applied + /// inside the client (at the raw byte layer, before SSE parsing) rather than + /// by the caller. + /// + /// Custom clients that parse SSE responses should override this method. + /// The default implementation delegates to [`Self::post_message`]. + fn post_message_with_max_sse_event_size( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_header: Option, + custom_headers: HashMap, + _max_sse_event_size: usize, + ) -> impl Future>> + + Send + + '_ { + self.post_message(uri, message, session_id, auth_header, custom_headers) + } fn delete_session( &self, uri: Arc, @@ -215,10 +450,14 @@ pub trait StreamableHttpClient: Clone + Send + 'static { auth_header: Option, custom_headers: HashMap, ) -> impl Future>> + Send + '_; + /// Open an SSE stream, optionally scoped to a legacy session. + /// + /// `session_id` is `None` when resuming a stateless response using only + /// `last_event_id`. fn get_stream( &self, uri: Arc, - session_id: Arc, + session_id: Option>, last_event_id: Option, auth_header: Option, custom_headers: HashMap, @@ -229,6 +468,33 @@ pub trait StreamableHttpClient: Clone + Send + 'static { >, > + Send + '_; + /// Open an SSE stream while enforcing a maximum raw event size. + /// + /// `max_sse_event_size` is not a per-request option: it is the + /// transport-wide [`StreamableHttpClientTransportConfig::max_sse_event_size`] + /// value, passed identically on every call because the limit must be applied + /// inside the client (at the raw byte layer, before SSE parsing) rather than + /// by the caller. + /// + /// Custom clients that parse SSE responses should override this method. + /// The default implementation delegates to [`Self::get_stream`]. + fn get_stream_with_max_sse_event_size( + &self, + uri: Arc, + session_id: Option>, + last_event_id: Option, + auth_header: Option, + custom_headers: HashMap, + _max_sse_event_size: usize, + ) -> impl Future< + Output = Result< + BoxStream<'static, Result>, + StreamableHttpError, + >, + > + Send + + '_ { + self.get_stream(uri, session_id, last_event_id, auth_header, custom_headers) + } } #[non_exhaustive] @@ -239,10 +505,11 @@ pub struct RetryConfig { struct StreamableHttpClientReconnect { pub client: C, - pub session_id: Arc, + pub session_id: Option>, pub uri: Arc, pub auth_header: Option, pub custom_headers: HashMap, + pub max_sse_event_size: usize, } impl SseStreamReconnect for StreamableHttpClientReconnect { @@ -254,13 +521,25 @@ impl SseStreamReconnect for StreamableHttpClientReconne let session_id = self.session_id.clone(); let auth_header = self.auth_header.clone(); let custom_headers = self.custom_headers.clone(); + let max_sse_event_size = self.max_sse_event_size; let last_event_id = last_event_id.map(|s| s.to_owned()); Box::pin(async move { client - .get_stream(uri, session_id, last_event_id, auth_header, custom_headers) + .get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + auth_header, + custom_headers, + max_sse_event_size, + ) .await }) } + + fn map_fatal_stream_error(&mut self, error: SseError) -> Option { + Some(StreamableHttpError::Sse(error)) + } } /// Info retained for cleaning up the session when the worker exits. @@ -279,6 +558,21 @@ pub struct StreamableHttpClientWorker { pub config: StreamableHttpClientTransportConfig, } +struct PostResult { + send_request: WorkerSendRequest>, + // None means the send future was dropped, or the request or transport was cancelled. + response: Option>>, + // The protocol version used to send this POST. + version: ProtocolVersion, +} + +struct PostSession { + id: Option>, + headers: HashMap, + version: ProtocolVersion, + cancellation: CancellationToken, +} + impl StreamableHttpClientWorker { pub fn new_simple(url: impl Into>) -> Self { Self { @@ -298,42 +592,279 @@ impl StreamableHttpClientWorker { } impl StreamableHttpClientWorker { - /// Convert a raw SSE stream into a JSON-RPC message stream without - /// reconnection logic. - fn raw_sse_to_jsonrpc( - stream: BoxedSseStream, - ) -> impl Stream>> + Send + 'static - { - stream.filter_map(|event| async { - match event { - Err(e) => Some(Err(StreamableHttpError::Sse(e))), - Ok(sse) => { - let is_message = - matches!(sse.event.as_deref(), None | Some("") | Some("message")); - if !is_message { - return None; - } - let data = sse.data?; - if data.trim().is_empty() { - return None; - } - match serde_json::from_str::(&data) { - Ok(msg) => Some(Ok(msg)), - Err(e) => { - tracing::debug!("failed to deserialize server message: {e}"); - None - } - } - } + // Run initialization and protocol-version changes without other ordinary POSTs. + fn is_ordering_barrier( + message: &ClientJsonRpcMessage, + negotiated_version: &ProtocolVersion, + ) -> bool { + match message { + ClientJsonRpcMessage::Request(request) => { + matches!( + &request.request, + ClientRequest::InitializeRequest(_) | ClientRequest::DiscoverRequest(_) + ) || request + .request + .get_meta() + .protocol_version() + .is_some_and(|version| &version != negotiated_version) + } + ClientJsonRpcMessage::Notification(notification) => matches!( + ¬ification.notification, + ClientNotification::InitializedNotification(_) + ), + _ => false, + } + } + + fn post_request( + client: C, + config: &StreamableHttpClientTransportConfig, + mut send_request: WorkerSendRequest, + session: PostSession, + transport_cancellation: CancellationToken, + ) -> BoxFuture<'static, PostResult> { + let uri = config.uri.clone(); + let auth_header = config.auth_header.clone(); + let max_sse_event_size = config.max_sse_event_size; + let control_request_timeout = config.control_request_timeout; + let is_control = Self::is_control_message(&send_request.message); + let cancellation = send_request + .cancellation_token() + .unwrap_or_else(|| transport_cancellation.child_token()); + Box::pin(async move { + let response = tokio::select! { + biased; + _ = cancellation.cancelled() => None, + _ = send_request.responder.closed() => None, + _ = session.cancellation.cancelled() => { + Some(Err(StreamableHttpError::SessionRecoveryTimeout)) + }, + _ = tokio::time::sleep(control_request_timeout), if is_control => { + Some(Err(StreamableHttpError::ControlRequestTimeout)) + }, + response = client.post_message_with_max_sse_event_size( + uri, + send_request.message.clone(), + session.id, + auth_header, + session.headers, + max_sse_event_size, + ) => Some(response), + }; + PostResult { + send_request, + response, + version: session.version, } }) } + fn cancellation_request_id(message: &ClientJsonRpcMessage) -> Option<&RequestId> { + match message { + ClientJsonRpcMessage::Notification(notification) => match ¬ification.notification { + ClientNotification::CancelledNotification(cancelled) => { + cancelled.params.request_id.as_ref() + } + _ => None, + }, + _ => None, + } + } + + fn client_request_id(message: &ClientJsonRpcMessage) -> Option { + match message { + ClientJsonRpcMessage::Request(request) => Some(request.id.clone()), + _ => None, + } + } + + fn server_response_id(message: &ServerJsonRpcMessage) -> Option<&RequestId> { + match message { + ServerJsonRpcMessage::Response(response) => Some(&response.id), + ServerJsonRpcMessage::Error(error) => error.id.as_ref(), + _ => None, + } + } + + fn clear_stream_response_pending( + pending_stream_response_ids: &mut HashSet, + message: &ServerJsonRpcMessage, + ) -> Option { + let response_id = Self::server_response_id(message)?; + if let Some(id) = pending_stream_response_ids.take(response_id) { + return Some(id); + } + let id = RequestId::Number(response_id.numeric_string_value()?); + pending_stream_response_ids.take(&id) + } + + async fn drain_queued_stream_messages( + sse_worker_rx: &mut tokio::sync::mpsc::Receiver, + context: &mut super::worker::WorkerContext, + pending_stream_response_ids: &mut HashSet, + ) -> Result<(), WorkerQuitReason>> { + loop { + match sse_worker_rx.try_recv() { + Ok(message) => { + let _ = + Self::clear_stream_response_pending(pending_stream_response_ids, &message); + context.send_to_handler(message).await?; + } + Err(tokio::sync::mpsc::error::TryRecvError::Empty) => return Ok(()), + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => return Ok(()), + } + } + } + + async fn fail_pending_stream_responses( + context: &mut super::worker::WorkerContext, + pending_stream_response_ids: &mut HashSet, + ) -> Result<(), WorkerQuitReason>> { + if pending_stream_response_ids.is_empty() { + return Ok(()); + } + + let pending_ids = std::mem::take(pending_stream_response_ids); + for id in pending_ids { + context + .send_to_handler(ServerJsonRpcMessage::error( + ErrorData::internal_error( + "streamable HTTP session was re-initialized before the response arrived", + None, + ), + Some(id), + )) + .await?; + } + Ok(()) + } + + async fn fail_pending_responses_except_retries( + context: &mut super::worker::WorkerContext, + pending_stream_response_ids: &mut HashSet, + recovery_posts: &VecDeque>, + ) -> Result<(), WorkerQuitReason>> { + // Keep only retries that have not already received a stream response. + let retry_ids: Vec<_> = recovery_posts + .iter() + .filter_map(|request| Self::client_request_id(&request.message)) + .filter(|id| pending_stream_response_ids.remove(id)) + .collect(); + Self::fail_pending_stream_responses(context, pending_stream_response_ids).await?; + pending_stream_response_ids.extend(retry_ids); + Ok(()) + } + + fn fail_recovery_posts( + recovery_posts: &mut VecDeque>, + pending_stream_response_ids: &mut HashSet, + error: StreamableHttpError, + ) { + // The backend error cannot be cloned. Return it to one caller + // and return the original session-expired error to the others. + let mut recovery_error = Some(error); + for send_request in recovery_posts.drain(..) { + let pending = Self::client_request_id(&send_request.message) + .is_none_or(|id| pending_stream_response_ids.remove(&id)); + let result = if pending { + Err(recovery_error + .take() + .unwrap_or(StreamableHttpError::SessionExpired)) + } else { + Ok(()) + }; + let _ = send_request.responder.send(result); + } + } + + /// Convert an SSE stream into JSON-RPC messages with reconnect semantics. + /// + /// This is used for request-scoped SSE responses as well as the standalone + /// GET stream. A request-scoped stream can close before its response arrives, + /// and SEP-1699 requires the client to honor `retry` and resume with + /// `Last-Event-ID` in that case. + fn reconnecting_sse_to_jsonrpc( + stream: BoxedSseStream, + client: C, + session_id: Option>, + uri: Arc, + auth_header: Option, + custom_headers: HashMap, + max_sse_event_size: usize, + retry_config: Arc, + ) -> impl Stream>> + Send + 'static + { + SseAutoReconnectStream::new_after_event_id( + stream, + StreamableHttpClientReconnect { + client, + session_id, + uri, + auth_header, + custom_headers, + max_sse_event_size, + }, + retry_config, + ) + } + + /// Convert a POST response SSE stream into JSON-RPC messages. + /// + /// Request-scoped streams resume via GET once the server has supplied an + /// event ID. The session header remains optional for stateless transports. + fn response_sse_to_jsonrpc( + stream: BoxedSseStream, + session_id: Option>, + client: C, + uri: Arc, + auth_header: Option, + custom_headers: HashMap, + max_sse_event_size: usize, + retry_config: Arc, + ) -> BoxStream<'static, Result>> { + Self::reconnecting_sse_to_jsonrpc( + stream, + client, + session_id, + uri, + auth_header, + custom_headers, + max_sse_event_size, + retry_config, + ) + .boxed() + } + + async fn run_response_stream( + mut sse_stream: BoxStream< + 'static, + Result>, + >, + sse_worker_tx: tokio::sync::mpsc::Sender, + origin: InboundStreamOrigin, + request_ct: CancellationToken, + stream_ct: CancellationToken, + uses_modern_http: bool, + ) -> Result<(), StreamableHttpError> { + tokio::select! { + biased; + _ = request_ct.cancelled(), if !uses_modern_http => { + // Stop reading, but keep the stream until the adapter + // handles cancellation or the send is dropped. + stream_ct.cancelled().await; + Ok(()) + } + result = Self::execute_sse_stream( + sse_stream.as_mut(), sse_worker_tx, origin, true, stream_ct.clone(), + ) => result, + } + } + async fn execute_sse_stream( sse_stream: impl Stream>> - + Send - + 'static, + + Send, sse_worker_tx: tokio::sync::mpsc::Sender, + origin: InboundStreamOrigin, close_on_response: bool, ct: CancellationToken, ) -> Result<(), StreamableHttpError> { @@ -348,9 +879,14 @@ impl StreamableHttpClientWorker { break; } }; - let Some(message) = message.transpose()? else { + let Some(mut message) = message.transpose()? else { break; }; + // SEP-2260: mark inbound requests with the stream they arrived on + // for the client receive-side association check. + if let ServerJsonRpcMessage::Request(request) = &mut message { + request.request.extensions_mut().insert(origin.clone()); + } let is_response = matches!( message, ServerJsonRpcMessage::Response(_) | ServerJsonRpcMessage::Error(_) @@ -374,6 +910,69 @@ impl StreamableHttpClientWorker { Ok(()) } + fn spawn_common_stream( + streams: &mut tokio::task::JoinSet>, + client: C, + session_id: Arc, + config: &StreamableHttpClientTransportConfig, + protocol_headers: HashMap, + sse_worker_tx: tokio::sync::mpsc::Sender, + transport_task_ct: CancellationToken, + ) { + let uri = config.uri.clone(); + let auth_header = config.auth_header.clone(); + let retry_config = config.retry_config.clone(); + let reconnect_uri = config.uri.clone(); + let reconnect_auth_header = config.auth_header.clone(); + let max_sse_event_size = config.max_sse_event_size; + + streams.spawn(async move { + let result = match client + .get_stream_with_max_sse_event_size( + uri, + Some(session_id.clone()), + None, + auth_header, + protocol_headers.clone(), + max_sse_event_size, + ) + .await + { + Ok(stream) => { + let sse_stream = SseAutoReconnectStream::new( + stream, + StreamableHttpClientReconnect { + client, + session_id: Some(session_id), + uri: reconnect_uri, + auth_header: reconnect_auth_header, + custom_headers: protocol_headers, + max_sse_event_size, + }, + retry_config, + ); + Self::execute_sse_stream( + sse_stream, + sse_worker_tx, + InboundStreamOrigin::Unassociated, + false, + transport_task_ct.child_token(), + ) + .await + } + Err(StreamableHttpError::ServerDoesNotSupportSse) => { + tracing::debug!("server doesn't support sse, skip common stream"); + Ok(()) + } + Err(error) => { + tracing::error!("fail to get common stream: {error}"); + Err(error) + } + }; + (None, result) + }); + } + /// Performs a transparent re-initialization handshake after a session-expired 404. /// /// Takes an owned clone of the client (avoiding `&self` across `.await` so the @@ -389,15 +988,23 @@ impl StreamableHttpClientWorker { uri: Arc, auth_header: Option, custom_headers: HashMap, - ) -> Result<(Option>, HashMap), StreamableHttpError> - { + max_sse_event_size: usize, + ) -> Result< + ( + Option>, + ProtocolVersion, + HashMap, + ), + StreamableHttpError, + > { let (init_msg, new_session_id_str) = client - .post_message( + .post_message_with_max_sse_event_size( uri.clone(), saved_init_request, None, auth_header.clone(), custom_headers.clone(), + max_sse_event_size, ) .await? .expect_initialized::() @@ -405,17 +1012,8 @@ impl StreamableHttpClientWorker { let new_session_id: Option> = new_session_id_str.map(|s| Arc::from(s.as_str())); - // Start from custom_headers, then inject the negotiated MCP-Protocol-Version - // so all subsequent requests carry the right version (MCP 2025-06-18 spec). - let mut new_protocol_headers = custom_headers; - if let ServerJsonRpcMessage::Response(response) = &init_msg { - if let ServerResult::InitializeResult(init_result) = &response.result { - if let Ok(hv) = HeaderValue::from_str(init_result.protocol_version.as_str()) { - new_protocol_headers - .insert(HeaderName::from_static("mcp-protocol-version"), hv); - } - } - } + let (negotiated_version, new_protocol_headers) = + negotiate_version_headers(&init_msg, custom_headers); let initialized_notification = ClientJsonRpcMessage::notification( ClientNotification::InitializedNotification(InitializedNotification { @@ -423,24 +1021,45 @@ impl StreamableHttpClientWorker { extensions: Default::default(), }), ); + // SEP-2243: notifications carry no Mcp-Param-*, so an empty tool cache suffices. + let initialized_headers = build_request_headers( + &new_protocol_headers, + &initialized_notification, + &HashMap::new(), + &negotiated_version, + ); client - .post_message( + .post_message_with_max_sse_event_size( uri, initialized_notification, new_session_id.clone(), auth_header, - new_protocol_headers.clone(), + initialized_headers, + max_sse_event_size, ) .await? .expect_accepted_or_json::()?; - Ok((new_session_id, new_protocol_headers)) + Ok((new_session_id, negotiated_version, new_protocol_headers)) } } impl Worker for StreamableHttpClientWorker { type Role = RoleClient; type Error = StreamableHttpError; + fn is_control_message(message: &ClientJsonRpcMessage) -> bool { + match message { + ClientJsonRpcMessage::Response(_) | ClientJsonRpcMessage::Error(_) => true, + ClientJsonRpcMessage::Notification(notification) => matches!( + notification.notification, + ClientNotification::CancelledNotification(_) + ), + ClientJsonRpcMessage::Request(_) => false, + } + } + fn supports_request_cancellation() -> bool { + true + } fn err_closed() -> Self::Error { StreamableHttpError::TransportChannelClosed } @@ -465,17 +1084,35 @@ impl Worker for StreamableHttpClientWorker { let _drop_guard = transport_task_ct.clone().drop_guard(); let WorkerSendRequest { responder, - message: initialize_request, + message: startup_request, + .. } = context.recv_from_handler().await?; - let saved_init_request = initialize_request.clone(); + let is_legacy_startup = matches!( + &startup_request, + ClientJsonRpcMessage::Request(request) + if matches!(&request.request, ClientRequest::InitializeRequest(_)) + ); + let mut saved_init_request = is_legacy_startup.then(|| startup_request.clone()); + let empty_tool_cache = HashMap::new(); + let (bootstrap_version, bootstrap_headers) = if is_legacy_startup { + (ProtocolVersion::default(), config.custom_headers.clone()) + } else { + request_version_headers( + &config.custom_headers, + &startup_request, + &ProtocolVersion::default(), + &empty_tool_cache, + ) + }; let (message, session_id) = match self .client - .post_message( + .post_message_with_max_sse_event_size( config.uri.clone(), - initialize_request, + startup_request, None, config.auth_header.clone(), - config.custom_headers.clone(), + bootstrap_headers.clone(), + config.max_sse_event_size, ) .await { @@ -494,7 +1131,10 @@ impl Worker for StreamableHttpClientWorker { )); } }; - let mut session_id: Option> = if let Some(session_id) = session_id { + let mut uses_modern_http = !is_legacy_startup; + let mut session_id: Option> = if uses_modern_http { + None + } else if let Some(session_id) = session_id { Some(session_id.into()) } else { if !self.config.allow_stateless { @@ -505,21 +1145,14 @@ impl Worker for StreamableHttpClientWorker { } None }; - // Extract the negotiated protocol version from the init response - // and build a custom headers map that includes MCP-Protocol-Version - // for all subsequent HTTP requests (per MCP 2025-06-18 spec). - let mut protocol_headers = { - let mut headers = config.custom_headers.clone(); - if let ServerJsonRpcMessage::Response(response) = &message { - if let ServerResult::InitializeResult(init_result) = &response.result { - if let Ok(hv) = HeaderValue::from_str(init_result.protocol_version.as_str()) { - // HeaderName::from_static requires lowercase - headers.insert(HeaderName::from_static("mcp-protocol-version"), hv); - } - } - } - headers + let (mut negotiated_version, mut protocol_headers) = if is_legacy_startup { + negotiate_version_headers(&message, config.custom_headers.clone()) + } else { + (bootstrap_version, bootstrap_headers) }; + // SEP-2243: tool input schemas (name -> schema) cached from tools/list responses, + // used to promote annotated tools/call arguments to Mcp-Param-* headers. + let mut tool_header_cache: HashMap> = HashMap::new(); // Store session info for cleanup when run() exits (not spawned, so cleanup completes before close() returns) let mut session_cleanup_info = session_id.as_ref().map(|sid| SessionCleanupInfo { @@ -531,100 +1164,227 @@ impl Worker for StreamableHttpClientWorker { }); context.send_to_handler(message).await?; - let initialized_notification = context.recv_from_handler().await?; - // expect a initialized response - self.client - .post_message( - config.uri.clone(), - initialized_notification.message, - session_id.clone(), - config.auth_header.clone(), - protocol_headers.clone(), - ) - .await - .map_err(WorkerQuitReason::fatal_context( - "send initialized notification", - ))? - .expect_accepted_or_json::() - .map_err(WorkerQuitReason::fatal_context( - "process initialized notification response", - ))?; - let _ = initialized_notification.responder.send(Ok(())); - #[allow(clippy::large_enum_variant)] - enum Event { - ClientMessage(WorkerSendRequest), + if is_legacy_startup { + let initialized_notification = context.recv_from_handler().await?; + let initialized_headers = build_request_headers( + &protocol_headers, + &initialized_notification.message, + &tool_header_cache, + &negotiated_version, + ); + self.client + .post_message_with_max_sse_event_size( + config.uri.clone(), + initialized_notification.message, + session_id.clone(), + config.auth_header.clone(), + initialized_headers, + config.max_sse_event_size, + ) + .await + .map_err(WorkerQuitReason::fatal_context( + "send initialized notification", + ))? + .expect_accepted_or_json::() + .map_err(WorkerQuitReason::fatal_context( + "process initialized notification response", + ))?; + let _ = initialized_notification.responder.send(Ok(())); + } + #[expect( + clippy::large_enum_variant, + reason = "the event is short-lived and boxing would add allocation in the event loop" + )] + enum Event { + ClientMessage(WorkerSendRequest>), + ControlMessage(WorkerSendRequest>), + StartPost(WorkerSendRequest>), + PostResult(PostResult), + RecoveryTimeout, ServerMessage(ServerJsonRpcMessage), - StreamResult(Result<(), StreamableHttpError>), + StreamResult { + request_id: Option, + result: Result<(), StreamableHttpError>, + }, } let mut streams = tokio::task::JoinSet::new(); + let mut pending_stream_response_ids = HashSet::new(); + let mut request_stream_cancellations = + HashMap::>::new(); + let mut posts = FuturesUnordered::>>::new(); + let mut control_posts = FuturesUnordered::>>::new(); + let mut session_cancellation = CancellationToken::new(); + let mut pending_message: Option> = None; + let mut recovery_posts = VecDeque::>::new(); + let mut recovery_deadline: Option = None; + let mut retrying_recovery = false; + let mut barrier_in_flight = false; + let max_concurrent_requests = config.max_concurrent_requests.max(1); + let mut awaiting_fallback_initialized = false; if let Some(session_id) = &session_id { - let client = self.client.clone(); - let uri = config.uri.clone(); - let session_id = session_id.clone(); - let auth_header = config.auth_header.clone(); - let retry_config = self.config.retry_config.clone(); - let sse_worker_tx = sse_worker_tx.clone(); - let transport_task_ct = transport_task_ct.clone(); - let config_uri = config.uri.clone(); - let config_auth_header = config.auth_header.clone(); - let spawn_headers = protocol_headers.clone(); - - streams.spawn(async move { - match client - .get_stream( - uri.clone(), - session_id.clone(), - None, - auth_header.clone(), - spawn_headers.clone(), - ) - .await - { - Ok(stream) => { - let sse_stream = SseAutoReconnectStream::new( - stream, - StreamableHttpClientReconnect { - client: client.clone(), - session_id: session_id.clone(), - uri: config_uri, - auth_header: config_auth_header, - custom_headers: spawn_headers, - }, - retry_config, - ); - Self::execute_sse_stream( - sse_stream, - sse_worker_tx, - false, - transport_task_ct.child_token(), + Self::spawn_common_stream( + &mut streams, + self.client.clone(), + session_id.clone(), + &config, + protocol_headers.clone(), + sse_worker_tx.clone(), + transport_task_ct.clone(), + ); + } + // Each POST uses the session and headers chosen when it starts. + // Only this loop updates the current session and protocol version. + let loop_result: Result<(), WorkerQuitReason> = 'main_loop: loop { + if retrying_recovery && recovery_posts.is_empty() && posts.is_empty() { + retrying_recovery = false; + } + if !retrying_recovery + && !recovery_posts.is_empty() + && posts.is_empty() + && control_posts.is_empty() + { + // Old POSTs have finished or reached the drain deadline. + // Retry only ordinary POSTs that returned SessionExpired, at most once each. + session_cancellation.cancel(); + recovery_deadline = None; + let recovery = tokio::select! { + _ = transport_task_ct.cancelled() => { + break 'main_loop Err(WorkerQuitReason::Cancelled); + } + result = tokio::time::timeout( + config.session_recovery_timeout, + Self::perform_reinitialization( + self.client.clone(), + saved_init_request.clone().expect("session recovery requires an initialize request"), + config.uri.clone(), + config.auth_header.clone(), + config.custom_headers.clone(), + config.max_sse_event_size, + ), + ) => result.unwrap_or(Err(StreamableHttpError::SessionRecoveryTimeout)), + }; + match recovery { + Ok((new_session_id, new_version, new_headers)) => { + streams.abort_all(); + while streams.join_next().await.is_some() {} + request_stream_cancellations.clear(); + Self::drain_queued_stream_messages( + &mut sse_worker_rx, + &mut context, + &mut pending_stream_response_ids, ) - .await + .await?; + Self::fail_pending_responses_except_retries( + &mut context, + &mut pending_stream_response_ids, + &recovery_posts, + ) + .await?; + session_id = new_session_id; + negotiated_version = new_version; + protocol_headers = new_headers; + session_cleanup_info = session_id.as_ref().map(|sid| SessionCleanupInfo { + client: self.client.clone(), + uri: config.uri.clone(), + session_id: sid.clone(), + auth_header: config.auth_header.clone(), + protocol_headers: protocol_headers.clone(), + }); + // Do not send controls queued during recovery to the new session. + context.advance_control_generation(); + session_cancellation = CancellationToken::new(); + if let Some(session_id) = &session_id { + Self::spawn_common_stream( + &mut streams, + self.client.clone(), + session_id.clone(), + &config, + protocol_headers.clone(), + sse_worker_tx.clone(), + transport_task_ct.clone(), + ); + } + retrying_recovery = true; } - Err(StreamableHttpError::ServerDoesNotSupportSse) => { - tracing::debug!("server doesn't support sse, skip common stream"); - Ok(()) - } - Err(e) => { - // fail to get common stream - tracing::error!("fail to get common stream: {e}"); - Err(e) + Err(error) => { + session_cancellation = CancellationToken::new(); + Self::fail_recovery_posts( + &mut recovery_posts, + &mut pending_stream_response_ids, + error, + ); } } + continue; + } + + let has_post_capacity = posts.len() < max_concurrent_requests; + let may_start = (retrying_recovery || recovery_posts.is_empty()) + && !barrier_in_flight + && has_post_capacity; + let may_receive = may_start && pending_message.is_none() && !retrying_recovery; + let queued = if retrying_recovery { + recovery_posts.front_mut() + } else { + pending_message.as_mut() + }; + let has_queued = queued.is_some(); + let can_process_queued = queued.as_ref().is_some_and(|request| { + let retry_completed = retrying_recovery + && Self::client_request_id(&request.message) + .is_some_and(|id| !pending_stream_response_ids.contains(&id)); + let ordering_satisfied = + !Self::is_ordering_barrier(&request.message, &negotiated_version) + || (posts.is_empty() && control_posts.is_empty()); + retry_completed || (may_start && ordering_satisfied) }); - } - // Main event loop - capture exit reason so we can do cleanup before returning - let loop_result: Result<(), WorkerQuitReason> = 'main_loop: loop { let event = tokio::select! { + _ = async { + if can_process_queued { + return; + } + let request = queued.expect("a POST is queued"); + let cancellation = request.cancellation_token().unwrap_or_default(); + tokio::select! { + _ = request.responder.closed() => {} + _ = cancellation.cancelled() => {} + } + }, if has_queued => { + let request = if retrying_recovery { + recovery_posts.pop_front() + } else { + pending_message.take() + }; + Event::StartPost(request.expect("a POST is ready to start")) + } _ = transport_task_ct.cancelled() => { tracing::debug!("cancelled"); break 'main_loop Err(WorkerQuitReason::Cancelled); } - message = context.recv_from_handler() => { + message = context.from_handler_rx.recv(), if may_receive => { match message { - Ok(msg) => Event::ClientMessage(msg), - Err(e) => break 'main_loop Err(e), + Some(msg) => Event::ClientMessage(msg), + None => break 'main_loop Err(WorkerQuitReason::HandlerTerminated), } }, + message = context.control_from_handler_rx.recv(), + if control_posts.is_empty() && !session_cancellation.is_cancelled() => { + match message { + Some(msg) => Event::ControlMessage(msg), + None => break 'main_loop Err(WorkerQuitReason::HandlerTerminated), + } + }, + Some(result) = posts.next(), if !posts.is_empty() => { + Event::PostResult(result) + }, + Some(result) = control_posts.next(), if !control_posts.is_empty() => { + Event::PostResult(result) + }, + _ = async { + if let Some(deadline) = recovery_deadline { + tokio::time::sleep_until(deadline).await; + } + }, if recovery_deadline.is_some() => Event::RecoveryTimeout, message = sse_worker_rx.recv() => { let Some(message) = message else { tracing::trace!("transport dropped, exiting"); @@ -634,191 +1394,391 @@ impl Worker for StreamableHttpClientWorker { }, terminated_stream = streams.join_next(), if !streams.is_empty() => { match terminated_stream { - Some(result) => { - Event::StreamResult(result.map_err(StreamableHttpError::TokioJoinError).and_then(std::convert::identity)) - } - None => { - continue + Some(Ok((request_id, result))) => { + Event::StreamResult { request_id, result } } + Some(Err(error)) => Event::StreamResult { + request_id: None, + result: Err(StreamableHttpError::TokioJoinError(error)), + }, + None => continue, } } }; match event { Event::ClientMessage(send_request) => { - let WorkerSendRequest { message, responder } = send_request; - // Pass a clone to the first attempt so `message` is retained for a - // potential re-init retry. `post_message` takes ownership and the - // trait cannot be changed, so the clone is unavoidable. - let response = self - .client - .post_message( - config.uri.clone(), - message.clone(), - session_id.clone(), - config.auth_header.clone(), - protocol_headers.clone(), - ) - .await; - let send_result = match response { - Err(StreamableHttpError::SessionExpired) => { - if !config.reinit_on_expired_session { - Err(StreamableHttpError::SessionExpired) - } else { - // The server discarded the session (HTTP 404). Perform a - // fresh handshake once and replay the original message. - tracing::info!( - "session expired (HTTP 404), attempting transparent re-initialization" - ); - match Self::perform_reinitialization( - self.client.clone(), - saved_init_request.clone(), - config.uri.clone(), - config.auth_header.clone(), - config.custom_headers.clone(), - ) - .await - { - Ok((new_session_id, new_protocol_headers)) => { - // Old streams hold the stale session ID; abort them - // so the new standalone SSE stream takes over. - streams.abort_all(); - - session_id = new_session_id; - protocol_headers = new_protocol_headers; - session_cleanup_info = - session_id.as_ref().map(|sid| SessionCleanupInfo { - client: self.client.clone(), - uri: config.uri.clone(), - session_id: sid.clone(), - auth_header: config.auth_header.clone(), - protocol_headers: protocol_headers.clone(), - }); - - if let Some(new_sid) = &session_id { - let client = self.client.clone(); - let uri = config.uri.clone(); - let new_sid = new_sid.clone(); - let auth_header = config.auth_header.clone(); - let retry_config = self.config.retry_config.clone(); - let sse_tx = sse_worker_tx.clone(); - let task_ct = transport_task_ct.clone(); - let config_uri = config.uri.clone(); - let config_auth = config.auth_header.clone(); - let spawn_headers = protocol_headers.clone(); - streams.spawn(async move { - match client - .get_stream( - uri, - new_sid.clone(), - None, - auth_header.clone(), - spawn_headers.clone(), - ) - .await - { - Ok(stream) => { - let sse_stream = SseAutoReconnectStream::new( - stream, - StreamableHttpClientReconnect { - client: client.clone(), - session_id: new_sid, - uri: config_uri, - auth_header: config_auth, - custom_headers: spawn_headers, - }, - retry_config, - ); - Self::execute_sse_stream( - sse_stream, - sse_tx, - false, - task_ct.child_token(), - ) - .await - } - Err(StreamableHttpError::ServerDoesNotSupportSse) => { - tracing::debug!( - "server doesn't support sse after re-init" - ); - Ok(()) - } - Err(e) => { - tracing::error!( - "fail to get common stream after re-init: {e}" - ); - Err(e) - } - } - }); - } - - let retry_response = self - .client - .post_message( - config.uri.clone(), - message, - session_id.clone(), - config.auth_header.clone(), - protocol_headers.clone(), - ) - .await; - match retry_response { - Err(e) => Err(e), - Ok(StreamableHttpPostResponse::Accepted) => { - tracing::trace!( - "client message accepted after re-init" - ); - Ok(()) - } - Ok(StreamableHttpPostResponse::Json(msg, ..)) => { - context.send_to_handler(msg).await?; - Ok(()) - } - Ok(StreamableHttpPostResponse::Sse(stream, ..)) => { - streams.spawn(Self::execute_sse_stream( - Self::raw_sse_to_jsonrpc(stream), - sse_worker_tx.clone(), - true, - transport_task_ct.child_token(), - )); - tracing::trace!("got new sse stream after re-init"); - Ok(()) - } - } - } - Err(reinit_err) => Err(reinit_err), - } - } // else enable_reinit_on_expired_session + pending_message = Some(send_request); + } + Event::ControlMessage(send_request) => { + if send_request.responder.is_closed() { + continue; + } + let cancellation_request_id = + Self::cancellation_request_id(&send_request.message); + let stale = send_request.control_generation() != context.control_generation(); + if stale { + // Do not send old controls to a replacement session. + let result = match cancellation_request_id { + Some(_) => Ok(()), + None => Err(StreamableHttpError::SessionExpired), + }; + let _ = send_request.responder.send(result); + continue; + } + if let Some(request_id) = cancellation_request_id { + drop(request_stream_cancellations.remove(request_id)); + pending_stream_response_ids.remove(request_id); + if uses_modern_http { + let _ = send_request.responder.send(Ok(())); + continue; } + } + let (version, headers) = request_version_headers( + &protocol_headers, + &send_request.message, + &negotiated_version, + &tool_header_cache, + ); + control_posts.push(Self::post_request( + self.client.clone(), + &config, + send_request, + PostSession { + id: session_id.clone(), + headers, + version, + cancellation: session_cancellation.clone(), + }, + transport_task_ct.clone(), + )); + } + Event::RecoveryTimeout => { + recovery_deadline = None; + session_cancellation.cancel(); + tracing::warn!("old-session POSTs did not finish before the recovery deadline"); + } + Event::StartPost(send_request) => { + let request_id = Self::client_request_id(&send_request.message); + let send_cancelled = send_request.responder.is_closed() + || send_request + .cancellation_token() + .is_some_and(|token| token.is_cancelled()); + let retry_completed = retrying_recovery + && request_id + .as_ref() + .is_some_and(|id| !pending_stream_response_ids.contains(id)); + if send_cancelled || retry_completed { + if retrying_recovery && let Some(id) = &request_id { + pending_stream_response_ids.remove(id); + } + let _ = send_request.responder.send(Ok(())); + continue; + } + let message = &send_request.message; + let is_fallback_initialize = saved_init_request.is_none() + && matches!( + message, + ClientJsonRpcMessage::Request(request) + if matches!( + &request.request, + ClientRequest::InitializeRequest(_) + ) + ); + if is_fallback_initialize { + saved_init_request = Some(message.clone()); + let WorkerSendRequest { + message, responder, .. + } = send_request; + // Servers do not assign sessions to `server/discover`, so a + // fallback initialize starts from a clean slate: no session + // ID, no cleanup state, and no streams to tear down. + debug_assert!( + session_id.is_none() + && session_cleanup_info.is_none() + && streams.is_empty(), + "discover bootstrap must not create session state" + ); + uses_modern_http = false; + + let response = self + .client + .post_message_with_max_sse_event_size( + config.uri.clone(), + message, + None, + config.auth_header.clone(), + config.custom_headers.clone(), + config.max_sse_event_size, + ) + .await; + let response = match response { + Ok(response) => { + let _ = responder.send(Ok(())); + response + } + Err(error) => { + let _ = responder.send(Err(error)); + continue; + } + }; + let (initialize_response, new_session_id) = response + .expect_initialized::() + .await + .map_err(WorkerQuitReason::fatal_context( + "process fallback initialize response", + ))?; + session_id = new_session_id.map(Arc::from); + if session_id.is_none() && !config.allow_stateless { + return Err(WorkerQuitReason::fatal( + StreamableHttpError::::MissingSessionIdInResponse, + "process fallback initialize response", + )); + } + (negotiated_version, protocol_headers) = negotiate_version_headers( + &initialize_response, + config.custom_headers.clone(), + ); + session_cleanup_info = + session_id.as_ref().map(|session_id| SessionCleanupInfo { + client: self.client.clone(), + uri: config.uri.clone(), + session_id: session_id.clone(), + auth_header: config.auth_header.clone(), + protocol_headers: protocol_headers.clone(), + }); + context.send_to_handler(initialize_response).await?; + awaiting_fallback_initialized = true; + continue; + } + + let barrier = Self::is_ordering_barrier(message, &negotiated_version); + debug_assert!(!barrier || (posts.is_empty() && control_posts.is_empty())); + let inline_version = match message { + ClientJsonRpcMessage::Request(request) => { + request.request.get_meta().protocol_version() + } + _ => None, + }; + let (request_version, request_headers) = request_version_headers( + &protocol_headers, + message, + &negotiated_version, + &tool_header_cache, + ); + if inline_version.is_some() { + negotiated_version = request_version.clone(); + if let Ok(value) = HeaderValue::from_str(request_version.as_str()) { + protocol_headers + .insert(HeaderName::from_static("mcp-protocol-version"), value); + } + if let Some(cleanup) = &mut session_cleanup_info { + cleanup.protocol_headers = protocol_headers.clone(); + } + } + barrier_in_flight = barrier; + // The common stream can return a response before this POST finishes. + if let Some(request_id) = request_id { + pending_stream_response_ids.insert(request_id); + } + posts.push(Self::post_request( + self.client.clone(), + &config, + send_request, + PostSession { + id: session_id.clone(), + headers: request_headers, + version: request_version, + cancellation: session_cancellation.clone(), + }, + transport_task_ct.clone(), + )); + } + Event::PostResult(PostResult { + send_request, + response, + version, + }) => { + let is_control = Self::is_control_message(&send_request.message); + if !is_control { + // An ordering barrier runs without other ordinary POSTs. + barrier_in_flight = false; + } + let request_id = Self::client_request_id(&send_request.message); + if request_id + .as_ref() + .is_some_and(|id| !pending_stream_response_ids.contains(id)) + { + // A stream response or cancellation already completed this request. + let _ = send_request.responder.send(Ok(())); + continue; + } + let will_retry = + matches!(&response, Some(Err(StreamableHttpError::SessionExpired))) + && !is_control + && !retrying_recovery + && config.reinit_on_expired_session + && saved_init_request.is_some(); + let awaits_stream_response = matches!( + &response, + Some(Ok(StreamableHttpPostResponse::Accepted + | StreamableHttpPostResponse::Sse(..))) + ); + if !will_retry + && !awaits_stream_response + && let Some(id) = &request_id + { + pending_stream_response_ids.remove(id); + } + let Some(response) = response else { + let _ = send_request.responder.send(Ok(())); + continue; + }; + if will_retry { + if recovery_posts.is_empty() { + recovery_deadline = + Some(tokio::time::Instant::now() + config.session_recovery_timeout); + } + recovery_posts.push_back(send_request); + continue; + } + let request_cancellation = send_request.cancellation_registration(); + let WorkerSendRequest { + message, responder, .. + } = send_request; + let is_initialized_notification = matches!( + &message, + ClientJsonRpcMessage::Notification(notification) + if matches!( + ¬ification.notification, + ClientNotification::InitializedNotification(_) + ) + ); + let send_result = match response { Err(e) => Err(e), Ok(StreamableHttpPostResponse::Accepted) => { tracing::trace!("client message accepted"); Ok(()) } - Ok(StreamableHttpPostResponse::Json(message, ..)) => { + Ok(StreamableHttpPostResponse::Json(mut message, ..)) => { + cache_tools_from_response( + &mut tool_header_cache, + &mut message, + &version, + ); context.send_to_handler(message).await?; Ok(()) } Ok(StreamableHttpPostResponse::Sse(stream, ..)) => { - streams.spawn(Self::execute_sse_stream( - Self::raw_sse_to_jsonrpc(stream), - sse_worker_tx.clone(), - true, - transport_task_ct.child_token(), - )); + let stream_request_id = request_id; + let sse_stream = Self::response_sse_to_jsonrpc( + stream, + session_id.clone(), + self.client.clone(), + config.uri.clone(), + config.auth_header.clone(), + protocol_headers.clone(), + config.max_sse_event_size, + self.config.retry_config.clone(), + ); + let request_ct = request_cancellation + .as_ref() + .map(|registration| registration.token()) + .unwrap_or_else(|| transport_task_ct.child_token()); + // A legacy adapter may need the open stream to handle cancellation. + let stream_ct = if uses_modern_http { + request_ct.clone() + } else { + request_cancellation + .as_ref() + .map(|registration| registration.lifetime_token()) + .unwrap_or_else(|| request_ct.clone()) + }; + if let (Some(request_id), Some(registration)) = + (stream_request_id.as_ref(), request_cancellation) + { + request_stream_cancellations + .insert(request_id.clone(), registration); + } + let stream_tx = sse_worker_tx.clone(); + let origin = match &stream_request_id { + Some(id) => InboundStreamOrigin::OutboundRequest(id.clone()), + None => InboundStreamOrigin::Unassociated, + }; + streams.spawn(async move { + let result = Self::run_response_stream( + sse_stream, + stream_tx, + origin, + request_ct, + stream_ct, + uses_modern_http, + ) + .await; + (stream_request_id, result) + }); tracing::trace!("got new sse stream"); Ok(()) } }; + if send_result.is_ok() + && awaiting_fallback_initialized + && is_initialized_notification + { + if let Some(session_id) = &session_id { + Self::spawn_common_stream( + &mut streams, + self.client.clone(), + session_id.clone(), + &config, + protocol_headers.clone(), + sse_worker_tx.clone(), + transport_task_ct.clone(), + ); + } + awaiting_fallback_initialized = false; + } let _ = responder.send(send_result); } - Event::ServerMessage(json_rpc_message) => { + Event::ServerMessage(mut json_rpc_message) => { + // Match against all pending requests, not just open response streams. + if let Some(request_id) = Self::clear_stream_response_pending( + &mut pending_stream_response_ids, + &json_rpc_message, + ) { + drop(request_stream_cancellations.remove(&request_id)); + } + cache_tools_from_response( + &mut tool_header_cache, + &mut json_rpc_message, + &negotiated_version, + ); // send the message to the handler if let Err(e) = context.send_to_handler(json_rpc_message).await { break 'main_loop Err(e); } } - Event::StreamResult(result) => { + Event::StreamResult { request_id, result } => { + if let Some(request_id) = request_id { + Self::drain_queued_stream_messages( + &mut sse_worker_rx, + &mut context, + &mut pending_stream_response_ids, + ) + .await?; + let cancelled = request_stream_cancellations + .remove(&request_id) + .is_some_and(|registration| registration.token().is_cancelled()); + if pending_stream_response_ids.remove(&request_id) && !cancelled { + context + .send_to_handler(ServerJsonRpcMessage::error( + ErrorData::transport_closed( + "streamable HTTP response stream closed before its final response", + ), + Some(request_id), + )) + .await?; + } + } if result.is_err() { tracing::warn!( "sse client event stream terminated with error: {:?}", @@ -829,10 +1789,17 @@ impl Worker for StreamableHttpClientWorker { } }; + // Stop outstanding http requests before deleting their session. + transport_task_ct.cancel(); + drop(posts); + drop(control_posts); + drop(pending_message); + drop(recovery_posts); + streams.abort_all(); + // Cleanup session before returning (ensures close() waits for session deletion) // Use a timeout to prevent indefinite hangs if the server is unresponsive if let Some(cleanup) = session_cleanup_info { - const SESSION_CLEANUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); let cleanup_session_id = cleanup.session_id.clone(); match tokio::time::timeout( SESSION_CLEANUP_TIMEOUT, @@ -950,7 +1917,7 @@ impl Worker for StreamableHttpClientWorker { /// async fn get_stream( /// &self, /// _uri: Arc, -/// _session_id: Arc, +/// _session_id: Option>, /// _last_event_id: Option, /// _auth_header: Option, /// _custom_headers: HashMap, @@ -1037,7 +2004,7 @@ impl StreamableHttpClientTransport { /// async fn get_stream( /// &self, /// _uri: Arc, - /// _session_id: Arc, + /// _session_id: Option>, /// _last_event_id: Option, /// _auth_header: Option, /// _custom_headers: HashMap, @@ -1062,21 +2029,42 @@ pub struct StreamableHttpClientTransportConfig { pub uri: Arc, pub retry_config: Arc, pub channel_buffer_capacity: usize, + /// Maximum number of ordinary http POSTs in progress (default: 16). + /// A POST stops counting when it completes or opens an sse response stream. + /// Zero is treated as one. Cancellation and replies use a separate queue + /// with one extra POST slot and use [`Self::control_request_timeout`]. + pub max_concurrent_requests: usize, + /// Maximum time a cancellation or reply POST can run after it starts (default: five seconds). + /// Time spent waiting in the control queue does not count toward this timeout. + pub control_request_timeout: Duration, + /// Maximum wait for old POSTs to finish before session recovery (default: five seconds). + /// The new initialization handshake has a separate timeout of the same length. + /// An unfinished old POST returns [`StreamableHttpError::SessionRecoveryTimeout`] + /// and is not retried because the server may have processed it. + pub session_recovery_timeout: Duration, /// if true, the transport will not require a session to be established pub allow_stateless: bool, /// The value to send in the authorization header pub auth_header: Option, /// Custom HTTP headers to include with every request pub custom_headers: HashMap, - /// Enables transparent recovery when the server reports an expired session (`HTTP 404`). + /// Maximum raw size of one SSE event accepted from the server. /// - /// When enabled, the transport performs one automatic recovery attempt: - /// 1. Replays the original `initialize` handshake to create a new session. - /// 2. Re-establishes streaming state for that session. - /// 3. Retries the in-flight request that failed with `SessionExpired`. + /// The built-in reqwest and Unix socket clients enforce this value. Custom + /// [`StreamableHttpClient`] implementations must override the corresponding + /// `*_with_max_sse_event_size` methods to enforce it. + pub max_sse_event_size: usize, + /// Automatically creates a new session when the server reports an expired + /// session (`http 404`). /// - /// This recovery is best-effort and bounded to a single attempt. If recovery fails, - /// the original failure path is preserved and the error is returned to the caller. + /// Ordinary POSTs that fail with `SessionExpired` in the same session share one + /// recovery attempt: + /// 1. Wait for old POSTs, up to [`Self::session_recovery_timeout`]. + /// 2. Repeat the original `initialize` handshake and open new streams. + /// 3. Retry each ordinary POST that failed with `SessionExpired` once. + /// + /// Control POSTs and other POST failures are not retried. If recovery or a retry + /// fails, the transport returns an error to the caller. pub reinit_on_expired_session: bool, } @@ -1088,6 +2076,24 @@ impl StreamableHttpClientTransportConfig { } } + /// Set how many ordinary POSTs can run at once. One keeps them serial; zero also means one. + pub fn max_concurrent_requests(mut self, limit: usize) -> Self { + self.max_concurrent_requests = limit.max(1); + self + } + + /// Set the timeout for cancellation and reply POSTs, starting when each POST starts. + pub fn control_request_timeout(mut self, timeout: Duration) -> Self { + self.control_request_timeout = timeout; + self + } + + /// Set the separate timeouts for waiting for old POSTs and creating a replacement session. + pub fn session_recovery_timeout(mut self, timeout: Duration) -> Self { + self.session_recovery_timeout = timeout; + self + } + /// Set the authorization header to send with requests /// /// # Arguments @@ -1126,6 +2132,12 @@ impl StreamableHttpClientTransportConfig { self } + /// Set the maximum raw size of one SSE event accepted from the server. + pub fn max_sse_event_size(mut self, bytes: usize) -> Self { + self.max_sse_event_size = bytes; + self + } + /// Set whether the transport should attempt transparent re-initialization on session expiration /// See [`Self::reinit_on_expired_session`] for details. /// # Example @@ -1146,10 +2158,448 @@ impl Default for StreamableHttpClientTransportConfig { uri: "localhost".into(), retry_config: Arc::new(ExponentialBackoff::default()), channel_buffer_capacity: 16, + max_concurrent_requests: 16, + control_request_timeout: Duration::from_secs(5), + session_recovery_timeout: Duration::from_secs(5), allow_stateless: true, auth_header: None, custom_headers: HashMap::new(), + max_sse_event_size: DEFAULT_MAX_SSE_EVENT_SIZE, reinit_on_expired_session: true, } } } + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use serde_json::json; + + use super::*; + use crate::{ + model::{ + GetExtensions, ListToolsResult, NumberOrString, ServerRequest, ServerResult, Tool, + }, + service::InboundStreamOrigin, + }; + + #[test] + fn json_rpc_body_preview_empty_html_and_object() { + assert_eq!(json_rpc_body_preview(b""), ""); + assert_eq!(json_rpc_body_preview(b"{}"), "{}"); + assert_eq!( + json_rpc_body_preview(b"not json"), + "not json" + ); + } + + #[test] + fn json_rpc_body_preview_truncates_long_bodies() { + let long = vec![b'x'; JSON_RPC_BODY_PREVIEW_LEN + 40]; + let preview = json_rpc_body_preview(&long); + assert_eq!(preview.len(), JSON_RPC_BODY_PREVIEW_LEN); + assert!(preview.chars().all(|c| c == 'x')); + } + + #[test] + fn json_rpc_parse_error_message_includes_preview() { + let err = serde_json::from_slice::(b"{}").unwrap_err(); + let msg = json_rpc_parse_error_message(&err, b"{}"); + assert!(msg.contains("could not parse JSON response as ServerJsonRpcMessage")); + assert!(msg.contains("{}"), "missing body preview: {msg}"); + } + + #[expect( + deprecated, + reason = "Sampling is deprecated by SEP-2577 but remains the canonical restricted request" + )] + fn sampling_request_message(id: i64) -> ServerJsonRpcMessage { + use crate::model::{CreateMessageRequest, CreateMessageRequestParams, SamplingMessage}; + ServerJsonRpcMessage::request( + ServerRequest::CreateMessageRequest(CreateMessageRequest::new( + CreateMessageRequestParams::new(vec![SamplingMessage::user_text("hi")], 16), + )), + NumberOrString::Number(id), + ) + } + + #[tokio::test] + async fn execute_sse_stream_marks_inbound_requests_with_origin() { + for origin in [ + InboundStreamOrigin::Unassociated, + InboundStreamOrigin::OutboundRequest(RequestId::Number(3)), + ] { + let response = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::default()), + NumberOrString::Number(1), + ); + let stream = futures::stream::iter([Ok(sampling_request_message(9)), Ok(response)]); + let (tx, mut rx) = tokio::sync::mpsc::channel(4); + StreamableHttpClientWorker::::execute_sse_stream( + stream, + tx, + origin.clone(), + false, + CancellationToken::new(), + ) + .await + .expect("stream completes"); + + let ServerJsonRpcMessage::Request(request) = + rx.recv().await.expect("request forwarded") + else { + panic!("expected request first"); + }; + assert_eq!( + request.request.extensions().get::(), + Some(&origin), + "inbound requests must carry their stream origin" + ); + // Responses are correlated by JSON-RPC id; no marker needed or added. + assert!(matches!( + rx.recv().await.expect("response forwarded"), + ServerJsonRpcMessage::Response(_) + )); + } + } + + type ReconnectAttempt = (Option, Option); + + #[derive(Clone, Default)] + struct StatelessReconnectClient { + reconnects: Arc>>, + } + + impl StreamableHttpClient for StatelessReconnectClient { + type Error = std::io::Error; + + async fn post_message( + &self, + _uri: Arc, + _message: ClientJsonRpcMessage, + _session_id: Option>, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + Err(StreamableHttpError::UnexpectedServerResponse( + "unexpected POST".into(), + )) + } + + async fn delete_session( + &self, + _uri: Arc, + _session_id: Arc, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + Ok(()) + } + + async fn get_stream( + &self, + _uri: Arc, + session_id: Option>, + last_event_id: Option, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + self.reconnects + .lock() + .expect("lock reconnects") + .push((session_id.map(|id| id.to_string()), last_event_id)); + let response = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::default()), + NumberOrString::Number(1), + ); + Ok(futures::stream::once(async move { + Ok(Sse { + event: None, + data: Some(serde_json::to_string(&response).expect("serialize response")), + id: Some("event-1".into()), + retry: None, + }) + }) + .boxed()) + } + } + + #[tokio::test] + async fn stateless_response_reconnects_with_last_event_id() { + let initial = futures::stream::iter([Ok(Sse { + event: None, + data: None, + id: Some("event-0".into()), + retry: Some(0), + })]) + .boxed(); + let client = StatelessReconnectClient::default(); + let reconnects = client.reconnects.clone(); + let stream = + StreamableHttpClientWorker::::response_sse_to_jsonrpc( + initial, + None, + client, + Arc::from("http://localhost/mcp"), + None, + HashMap::new(), + DEFAULT_MAX_SSE_EVENT_SIZE, + Arc::new(ExponentialBackoff { + max_times: Some(1), + base_duration: Duration::ZERO, + max_delay: None, + }), + ); + let mut stream = std::pin::pin!(stream); + + let message = stream.next().await.expect("replayed response").unwrap(); + + assert!(matches!(message, ServerJsonRpcMessage::Response(_))); + assert_eq!( + reconnects.lock().expect("lock reconnects").as_slice(), + &[(None, Some("event-0".into()))] + ); + } + + #[derive(Clone, Default)] + struct ResumedRequestClient { + reconnects: Arc>>, + } + + impl StreamableHttpClient for ResumedRequestClient { + type Error = std::io::Error; + + async fn post_message( + &self, + _uri: Arc, + _message: ClientJsonRpcMessage, + _session_id: Option>, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + Err(StreamableHttpError::UnexpectedServerResponse( + "unexpected POST".into(), + )) + } + + async fn delete_session( + &self, + _uri: Arc, + _session_id: Arc, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + Ok(()) + } + + async fn get_stream( + &self, + _uri: Arc, + session_id: Option>, + last_event_id: Option, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + self.reconnects + .lock() + .expect("lock reconnects") + .push((session_id.map(|id| id.to_string()), last_event_id)); + let request = sampling_request_message(9); + let response = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::default()), + NumberOrString::Number(1), + ); + // Stay open after the response, like a live connection, so the + // post-response drain in `execute_sse_stream` doesn't trigger + // further reconnects. + Ok(futures::stream::iter([request, response].map(|message| { + Ok(Sse { + event: None, + data: Some(serde_json::to_string(&message).expect("serialize message")), + id: None, + retry: None, + }) + })) + .chain(futures::stream::pending()) + .boxed()) + } + } + + /// SEP-1699 resumes a broken POST SSE stream via GET + Last-Event-ID + /// beneath `execute_sse_stream`, so the SEP-2260 origin marker must span + /// resumes; if reconnection were hoisted above the marker attach point, + /// replayed associated requests would be wrongly rejected with -32602. + #[tokio::test] + async fn resumed_post_stream_requests_keep_outbound_origin() { + let initial = futures::stream::iter([Ok(Sse { + event: None, + data: None, + id: Some("e1".into()), + retry: Some(0), + })]) + .boxed(); + let client = ResumedRequestClient::default(); + let reconnects = client.reconnects.clone(); + let sse_stream = + StreamableHttpClientWorker::::response_sse_to_jsonrpc( + initial, + None, + client, + Arc::from("http://localhost/mcp"), + None, + HashMap::new(), + DEFAULT_MAX_SSE_EVENT_SIZE, + Arc::new(ExponentialBackoff { + max_times: Some(1), + base_duration: Duration::ZERO, + max_delay: None, + }), + ); + + let origin = InboundStreamOrigin::OutboundRequest(RequestId::Number(3)); + let (tx, mut rx) = tokio::sync::mpsc::channel(4); + StreamableHttpClientWorker::::execute_sse_stream( + sse_stream, + tx, + origin.clone(), + true, + CancellationToken::new(), + ) + .await + .expect("stream completes"); + + assert_eq!( + reconnects.lock().expect("lock reconnects").as_slice(), + &[(None, Some("e1".into()))], + "the request must arrive on the resumed connection" + ); + let ServerJsonRpcMessage::Request(request) = rx.recv().await.expect("request forwarded") + else { + panic!("expected request first"); + }; + assert_eq!( + request.request.extensions().get::(), + Some(&origin), + "origin marker must survive SSE resumption" + ); + assert!(matches!( + rx.recv().await.expect("response forwarded"), + ServerJsonRpcMessage::Response(_) + )); + } + + fn tool(name: &'static str, annotation: serde_json::Value) -> Tool { + let schema = json!({ + "type": "object", + "properties": { + "value": annotation, + }, + }); + Tool::new( + name, + name, + Arc::new(schema.as_object().expect("object schema").clone()), + ) + } + + #[test] + fn cache_tools_removes_invalid_header_annotations() { + let valid = tool( + "valid", + json!({ "type": "string", "x-mcp-header": "Value" }), + ); + let invalid = tool("invalid", json!({ "type": "string", "x-mcp-header": "" })); + let mut message = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::with_all_items(vec![valid, invalid])), + NumberOrString::Number(1), + ); + let mut cache = HashMap::new(); + + cache_tools_from_response(&mut cache, &mut message, &ProtocolVersion::V_2026_07_28); + + let ServerJsonRpcMessage::Response(response) = &mut message else { + panic!("expected tools/list response"); + }; + let ServerResult::ListToolsResult(result) = &mut response.result else { + panic!("expected tools/list result"); + }; + assert_eq!( + ( + result + .tools + .iter() + .map(|tool| tool.name.as_ref()) + .collect::>(), + cache.keys().map(String::as_str).collect::>(), + ), + (vec!["valid"], vec!["valid"]) + ); + } + + #[test] + fn cache_tools_preserves_pre_standard_header_results() { + let invalid = tool("legacy", json!({ "type": "string", "x-mcp-header": "" })); + let mut message = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::with_all_items(vec![invalid])), + NumberOrString::Number(1), + ); + let mut cache = HashMap::new(); + + cache_tools_from_response(&mut cache, &mut message, &ProtocolVersion::V_2025_11_25); + + let ServerJsonRpcMessage::Response(response) = message else { + panic!("expected tools/list response"); + }; + let ServerResult::ListToolsResult(result) = response.result else { + panic!("expected tools/list result"); + }; + assert_eq!( + result + .tools + .iter() + .map(|tool| tool.name.as_ref()) + .collect::>(), + vec!["legacy"] + ); + } + + #[cfg(feature = "transport-streamable-http-client-reqwest")] + #[test] + fn clear_stream_response_pending_accepts_stringified_numeric_id() { + let mut pending = HashSet::from([NumberOrString::Number(1)]); + let response = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::default()), + NumberOrString::String("1".into()), + ); + + let matched_id = + StreamableHttpClientWorker::::clear_stream_response_pending( + &mut pending, + &response, + ); + + assert_eq!(matched_id, Some(NumberOrString::Number(1))); + assert!(pending.is_empty()); + } + + #[cfg(feature = "transport-streamable-http-client-reqwest")] + #[test] + fn clear_stream_response_pending_prefers_exact_string_id() { + let string_id = NumberOrString::String("1".into()); + let mut pending = HashSet::from([NumberOrString::Number(1), string_id.clone()]); + let response = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::default()), + string_id.clone(), + ); + + let matched_id = + StreamableHttpClientWorker::::clear_stream_response_pending( + &mut pending, + &response, + ); + + assert_eq!(matched_id, Some(string_id)); + assert_eq!(pending, HashSet::from([NumberOrString::Number(1)])); + } +} diff --git a/crates/rmcp/src/transport/streamable_http_server/session.rs b/crates/rmcp/src/transport/streamable_http_server/session.rs index 4be2651..ea2e85b 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session.rs @@ -13,13 +13,15 @@ //! //! * [`local::LocalSessionManager`] — in-memory session store (default). //! * [`never::NeverSessionManager`] — rejects all session operations, used -//! when stateful mode is disabled. +//! when legacy session mode is disabled. //! //! # Custom session managers //! //! Implement the [`SessionManager`] trait to back sessions with a database, //! Redis, or any other external store. +use std::sync::Arc; + use futures::Stream; pub use crate::transport::common::server_side_http::{ServerSseMessage, SessionId}; @@ -32,7 +34,10 @@ pub mod local; pub mod never; pub mod store; -pub use store::{SessionState, SessionStore, SessionStoreError}; +pub use store::{ + EventId, EventStore, EventStoreError, EventStream, SessionState, SessionStore, + SessionStoreError, StreamId, +}; /// Extension marker inserted into the `initialize` request extensions during a /// session restore replay. Handlers can check for its presence to distinguish a @@ -72,6 +77,16 @@ pub enum RestoreOutcome { /// trait for every HTTP request that carries (or should carry) a session ID. /// /// See the [module-level docs](self) for background on sessions. +/// +/// # SEP-2260 request association +/// +/// Server-initiated requests issued while handling a client request carry an +/// [`OriginatingRequestId`](crate::service::OriginatingRequestId) marker in +/// their non-serialized `Extensions`; the bundled local session manager uses +/// it to deliver such requests on the originating request's SSE stream. +/// Implementations that serialize messages between processes lose the marker +/// and fall back to the standalone stream, violating SEP-2260 for 2026-07-28+ +/// clients; they need their own association mechanism. pub trait SessionManager: Send + Sync + 'static { type Error: std::error::Error + Send + 'static; type Transport: crate::transport::Transport; @@ -151,4 +166,9 @@ pub trait SessionManager: Send + Sync + 'static { ) -> impl Future, Self::Error>> + Send { futures::future::ready(Ok(RestoreOutcome::NotSupported)) } + + /// Return the shared event store used for resumable SSE streams. + fn event_store(&self) -> Option> { + None + } } diff --git a/crates/rmcp/src/transport/streamable_http_server/session/local.rs b/crates/rmcp/src/transport/streamable_http_server/session/local.rs index 7e28932..e03e5b7 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/local.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/local.rs @@ -1,6 +1,7 @@ use std::{ collections::{HashMap, HashSet, VecDeque}, num::ParseIntError, + sync::Arc, time::{Duration, Instant}, }; @@ -32,6 +33,15 @@ use crate::{ pub struct LocalSessionManager { pub sessions: tokio::sync::RwLock>, pub session_config: SessionConfig, + event_store: Option>, +} + +impl LocalSessionManager { + /// Configure this session manager to use a shared event store. + pub fn with_event_store(mut self, event_store: Arc) -> Self { + self.event_store = Some(event_store); + self + } } #[derive(Debug, Error)] @@ -49,7 +59,11 @@ impl SessionManager for LocalSessionManager { type Transport = WorkerTransport; async fn create_session(&self) -> Result<(SessionId, Self::Transport), Self::Error> { let id = session_id(); - let (handle, worker) = create_local_session(id.clone(), self.session_config.clone()); + let (handle, worker) = create_local_session_with_event_store( + id.clone(), + self.session_config.clone(), + self.event_store.clone(), + ); self.sessions.write().await.insert(id.clone(), handle); Ok((id, WorkerTransport::spawn(worker))) } @@ -95,15 +109,7 @@ impl SessionManager for LocalSessionManager { let receiver = handle.establish_request_wise_channel().await?; let http_request_id = receiver.http_request_id; handle.push_message(message, http_request_id).await?; - - let priming = self.session_config.sse_retry.map(|retry| { - let event_id = match http_request_id { - Some(id) => format!("0/{id}"), - None => "0".into(), - }; - ServerSseMessage::priming(event_id, retry) - }); - Ok(futures::stream::iter(priming).chain(ReceiverStream::new(receiver.inner))) + Ok(ReceiverStream::new(receiver.inner)) } async fn create_standalone_stream( @@ -123,12 +129,19 @@ impl SessionManager for LocalSessionManager { id: &SessionId, last_event_id: String, ) -> Result + Send + 'static, Self::Error> { + if let Some(event_store) = &self.event_store { + let stream = event_store + .replay_events_after(&last_event_id) + .await + .map_err(SessionError::EventStore)?; + return Ok(stream.left_stream()); + } let sessions = self.sessions.read().await; let handle = sessions .get(id) .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?; let receiver = handle.resume(last_event_id.parse()?).await?; - Ok(ReceiverStream::new(receiver.inner)) + Ok(ReceiverStream::new(receiver.inner).right_stream()) } async fn accept_message( @@ -153,10 +166,18 @@ impl SessionManager for LocalSessionManager { // A concurrent request already restored this session. return Ok(RestoreOutcome::AlreadyPresent); } - let (handle, worker) = create_local_session(id.clone(), self.session_config.clone()); + let (handle, worker) = create_local_session_with_event_store( + id.clone(), + self.session_config.clone(), + self.event_store.clone(), + ); sessions.insert(id, handle); Ok(RestoreOutcome::Restored(WorkerTransport::spawn(worker))) } + + fn event_store(&self) -> Option> { + self.event_store.clone() + } } /// `/request_id>` @@ -209,7 +230,9 @@ impl std::str::FromStr for EventId { } } -use super::{RestoreOutcome, ServerSseMessage, SessionManager}; +use super::{ + EventStore, EventStoreError, RestoreOutcome, ServerSseMessage, SessionManager, StreamId, +}; struct CachedTx { tx: Sender, @@ -217,6 +240,8 @@ struct CachedTx { http_request_id: Option, capacity: usize, starting_index: usize, + stream_id: StreamId, + event_store: Option>, } impl CachedTx { @@ -224,6 +249,8 @@ impl CachedTx { tx: Sender, http_request_id: Option, starting_index: usize, + stream_id: StreamId, + event_store: Option>, ) -> Self { Self { cache: VecDeque::with_capacity(tx.capacity()), @@ -231,10 +258,16 @@ impl CachedTx { tx, http_request_id, starting_index, + stream_id, + event_store, } } - fn new_common(tx: Sender) -> Self { - Self::new(tx, None, 0) + fn new_common( + tx: Sender, + session_id: &SessionId, + event_store: Option>, + ) -> Self { + Self::new(tx, None, 0, format!("{session_id}:common"), event_store) } fn next_event_id(&self) -> EventId { @@ -253,16 +286,31 @@ impl CachedTx { } } - async fn send(&mut self, message: ServerJsonRpcMessage) { - let event_id = self.next_event_id(); - let message = ServerSseMessage::new(event_id.to_string(), message); - self.cache_and_send(message).await; + async fn send(&mut self, message: ServerJsonRpcMessage) -> Result<(), SessionError> { + self.store_cache_and_send(ServerSseMessage::from_message(message)) + .await } - async fn send_priming(&mut self, retry: Duration) { - let event_id = self.next_event_id(); - let message = ServerSseMessage::priming(event_id.to_string(), retry); - self.cache_and_send(message).await; + async fn send_priming(&mut self, retry: Duration) -> Result<(), SessionError> { + self.store_cache_and_send(ServerSseMessage::retry(retry)) + .await + } + + async fn store_cache_and_send( + &mut self, + mut event: ServerSseMessage, + ) -> Result<(), SessionError> { + let event_id = if let Some(event_store) = &self.event_store { + event_store + .store_event(&self.stream_id, &event) + .await + .map_err(SessionError::EventStore)? + } else { + self.next_event_id().to_string() + }; + event.event_id = Some(event_id); + self.cache_and_send(event).await; + Ok(()) } async fn cache_and_send(&mut self, message: ServerSseMessage) { @@ -330,6 +378,7 @@ pub struct LocalSessionWorker { shadow_txs: Vec>, event_rx: Receiver, session_config: SessionConfig, + event_store: Option>, } impl LocalSessionWorker { @@ -353,6 +402,8 @@ pub enum SessionError { InvalidEventId, #[error("IO error: {0}")] Io(#[from] std::io::Error), + #[error("Event store error: {0}")] + EventStore(#[source] EventStoreError), } impl From for std::io::Error { @@ -428,11 +479,11 @@ impl LocalSessionWorker { &mut self, notification: &JsonRpcNotification, ) { - if let ClientNotification::CancelledNotification(n) = ¬ification.notification { - if let Some(request_id) = n.params.request_id.clone() { - let resource = ResourceKey::McpRequestId(request_id); - self.unregister_resource(&resource); - } + if let ClientNotification::CancelledNotification(n) = ¬ification.notification + && let Some(request_id) = n.params.request_id.clone() + { + let resource = ResourceKey::McpRequestId(request_id); + self.unregister_resource(&resource); } } fn evict_expired_channels(&mut self) { @@ -450,12 +501,21 @@ impl LocalSessionWorker { ) -> Result { let http_request_id = self.next_http_request_id(); let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity); - let starting_index = usize::from(self.session_config.sse_retry.is_some()); + let mut cached_tx = CachedTx::new( + tx, + Some(http_request_id), + 0, + uuid::Uuid::new_v4().to_string(), + self.event_store.clone(), + ); + if let Some(retry) = self.session_config.sse_retry { + cached_tx.send_priming(retry).await?; + } self.tx_router.insert( http_request_id, HttpRequestWise { resources: Default::default(), - tx: CachedTx::new(tx, Some(http_request_id), starting_index), + tx: cached_tx, completed_at: None, }, ); @@ -467,7 +527,36 @@ impl LocalSessionWorker { } fn resolve_outbound_channel(&self, message: &ServerJsonRpcMessage) -> OutboundChannel { match &message { - ServerJsonRpcMessage::Request(_) => OutboundChannel::Common, + // SEP-2260: requests carrying an OriginatingRequestId marker ride the + // originating request's SSE stream, never the standalone GET stream. + ServerJsonRpcMessage::Request(json_rpc_request) => { + use crate::model::GetExtensions; + let originating = json_rpc_request + .request + .extensions() + .get::(); + match originating { + Some(originating) => match self + .resource_router + .get(&ResourceKey::McpRequestId(originating.0.clone())) + { + Some(id) => OutboundChannel::RequestWise { + id: *id, + close: false, + }, + None => { + tracing::warn!( + originating_request_id = %originating.0, + "associated server request could not be routed to its \ + originating stream (request completed or association \ + lost); falling back to standalone stream" + ); + OutboundChannel::Common + } + }, + None => OutboundChannel::Common, + } + } ServerJsonRpcMessage::Notification(JsonRpcNotification { notification: ServerNotification::ProgressNotification(Notification { @@ -548,19 +637,17 @@ impl LocalSessionWorker { match outbound_channel { OutboundChannel::RequestWise { id, close } => { if let Some(request_wise) = self.tx_router.get_mut(&id) { - request_wise.tx.send(message).await; - if close { - if let Some(channel) = self.tx_router.remove(&id) { - for resource in channel.resources { - self.resource_router.remove(&resource); - } + request_wise.tx.send(message).await?; + if close && let Some(channel) = self.tx_router.remove(&id) { + for resource in channel.resources { + self.resource_router.remove(&resource); } } } else { return Err(SessionError::ChannelClosed(Some(id))); } } - OutboundChannel::Common => self.common.send(message).await, + OutboundChannel::Common => self.common.send(message).await?, } Ok(()) } @@ -593,10 +680,20 @@ impl LocalSessionWorker { inner: rx, }) } - None => self.resume_or_shadow_common(last_event_id.index).await, + None => { + self.resume_or_shadow_common(Some(last_event_id.index)) + .await + } } } + async fn establish_common_channel( + &mut self, + ) -> Result { + let last_event_index = self.event_store.is_none().then_some(0); + self.resume_or_shadow_common(last_event_index).await + } + /// Resume the common channel, or create a shadow stream if the primary is /// still active. /// @@ -611,7 +708,7 @@ impl LocalSessionWorker { /// killing each other by repeatedly replacing the common channel sender. async fn resume_or_shadow_common( &mut self, - last_event_index: usize, + last_event_index: Option, ) -> Result { let is_replacing_dead_primary = self.common.tx.is_closed(); let capacity = if is_replacing_dead_primary { @@ -624,9 +721,9 @@ impl LocalSessionWorker { // Primary common channel is dead — replace it. tracing::debug!("Replacing dead common channel with new primary"); self.common.tx = tx; - // Replay cached messages from where the client left off so - // server-initiated requests and notifications are not lost. - self.common.sync(last_event_index).await?; + if let Some(last_event_index) = last_event_index { + self.common.sync(last_event_index).await?; + } } else { // Primary common channel is still active. Create a shadow stream // that stays alive via SSE keep-alive but doesn't receive @@ -668,7 +765,7 @@ impl LocalSessionWorker { // Send priming event if retry interval is specified if let Some(interval) = retry_interval { - request_wise.tx.send_priming(interval).await; + request_wise.tx.send_priming(interval).await?; } // Close the stream by dropping the sender @@ -685,7 +782,7 @@ impl LocalSessionWorker { None => { // Send priming event if retry interval is specified if let Some(interval) = retry_interval { - self.common.send_priming(interval).await; + self.common.send_priming(interval).await?; } // Close the stream by dropping the sender @@ -732,6 +829,9 @@ pub enum SessionEvent { retry_interval: Option, responder: oneshot::Sender>, }, + EstablishCommonChannel { + responder: oneshot::Sender>, + }, } #[derive(Debug, Clone)] @@ -820,13 +920,7 @@ impl LocalSessionHandle { ) -> Result { let (tx, rx) = tokio::sync::oneshot::channel(); self.event_tx - .send(SessionEvent::Resume { - last_event_id: EventId { - http_request_id: None, - index: 0, - }, - responder: tx, - }) + .send(SessionEvent::EstablishCommonChannel { responder: tx }) .await .map_err(|_| SessionError::SessionServiceTerminated)?; rx.await @@ -938,9 +1032,6 @@ pub enum LocalSessionWorkerError { FailToSendInitializeRequest(SessionError), #[error("fail to handle message: {0}")] FailToHandleMessage(SessionError), - #[deprecated(note = "idle timeout now surfaces as WorkerQuitReason::IdleTimeout")] - #[error("keep alive timeout after {}ms", _0.as_millis())] - KeepAliveTimeout(Duration), #[error("init timeout after {}ms", _0.as_millis())] InitTimeout(Duration), #[error("Transport closed")] @@ -1036,7 +1127,9 @@ impl Worker for LocalSessionWorker { } }; match event { - InnerEvent::FromHandler(WorkerSendRequest { message, responder }) => { + InnerEvent::FromHandler(WorkerSendRequest { + message, responder, .. + }) => { // catch response let to_unregister = match &message { crate::model::JsonRpcMessage::Response(json_rpc_response) => { @@ -1085,6 +1178,10 @@ impl Worker for LocalSessionWorker { let handle_result = self.establish_request_wise_channel().await; let _ = responder.send(handle_result); } + InnerEvent::FromHttpService(SessionEvent::EstablishCommonChannel { responder }) => { + let handle_result = self.establish_common_channel().await; + let _ = responder.send(handle_result); + } InnerEvent::FromHttpService(SessionEvent::CloseRequestWiseChannel { id, responder, @@ -1179,11 +1276,19 @@ impl Default for SessionConfig { pub fn create_local_session( id: impl Into, config: SessionConfig, +) -> (LocalSessionHandle, LocalSessionWorker) { + create_local_session_with_event_store(id, config, None) +} + +fn create_local_session_with_event_store( + id: impl Into, + config: SessionConfig, + event_store: Option>, ) -> (LocalSessionHandle, LocalSessionWorker) { let id = id.into(); let (event_tx, event_rx) = tokio::sync::mpsc::channel(config.channel_capacity); let (common_tx, _) = tokio::sync::mpsc::channel(config.channel_capacity); - let common = CachedTx::new_common(common_tx); + let common = CachedTx::new_common(common_tx, &id, event_store.clone()); tracing::info!(session_id = ?id, "create new session"); let handle = LocalSessionHandle { event_tx, @@ -1198,6 +1303,74 @@ pub fn create_local_session( shadow_txs: Vec::new(), event_rx, session_config: config.clone(), + event_store, }; (handle, session_worker) } + +#[cfg(test)] +mod sep2260_routing_tests { + use super::*; + use crate::service::OriginatingRequestId; + + fn roots_request(originating: Option) -> ServerJsonRpcMessage { + #[expect( + deprecated, + reason = "roots is SEP-2577-deprecated; any restricted request works here" + )] + let mut request = crate::model::ListRootsRequest { + method: Default::default(), + extensions: Default::default(), + }; + if let Some(id) = originating { + request.extensions.insert(OriginatingRequestId(id)); + } + ServerJsonRpcMessage::request( + crate::model::ServerRequest::ListRootsRequest(request), + RequestId::Number(1000), + ) + } + + #[tokio::test] + async fn associated_server_request_routes_to_originating_stream() { + let (_handle, mut worker) = create_local_session("test-session", SessionConfig::default()); + let receiver = worker.establish_request_wise_channel().await.unwrap(); + let http_request_id = receiver.http_request_id.unwrap(); + let originating_id = RequestId::Number(7); + worker.register_resource( + ResourceKey::McpRequestId(originating_id.clone()), + http_request_id, + ); + + let channel = worker.resolve_outbound_channel(&roots_request(Some(originating_id))); + assert!( + matches!(channel, OutboundChannel::RequestWise { id, close: false } if id == http_request_id) + ); + } + + #[tokio::test] + async fn unassociated_server_request_routes_to_common_stream() { + let (_handle, mut worker) = create_local_session("test-session", SessionConfig::default()); + let _receiver = worker.establish_request_wise_channel().await.unwrap(); + let channel = worker.resolve_outbound_channel(&roots_request(None)); + assert!(matches!(channel, OutboundChannel::Common)); + } + + #[tokio::test] + async fn associated_request_for_completed_request_falls_back_to_common() { + // Originating request already unregistered (completion race / re-scoped + // task op): must fall back to Common, never a closed request-wise channel. + let (_handle, mut worker) = create_local_session("test-session", SessionConfig::default()); + let receiver = worker.establish_request_wise_channel().await.unwrap(); + let http_request_id = receiver.http_request_id.unwrap(); + let originating_id = RequestId::Number(7); + worker.register_resource( + ResourceKey::McpRequestId(originating_id.clone()), + http_request_id, + ); + worker.unregister_resource(&ResourceKey::McpRequestId(originating_id.clone())); + + let channel = worker.resolve_outbound_channel(&roots_request(Some(originating_id))); + assert!(matches!(channel, OutboundChannel::Common)); + } +} diff --git a/crates/rmcp/src/transport/streamable_http_server/session/never.rs b/crates/rmcp/src/transport/streamable_http_server/session/never.rs index a2f72d8..d83031c 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/never.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/never.rs @@ -1,7 +1,9 @@ +use std::sync::Arc; + use futures::Stream; use thiserror::Error; -use super::{ServerSseMessage, SessionId, SessionManager}; +use super::{EventStore, ServerSseMessage, SessionId, SessionManager}; use crate::{ RoleServer, model::{ClientJsonRpcMessage, ServerJsonRpcMessage}, @@ -14,7 +16,17 @@ use crate::{ pub struct ErrorSessionManagementNotSupported; #[derive(Debug, Clone, Default)] #[non_exhaustive] -pub struct NeverSessionManager {} +pub struct NeverSessionManager { + event_store: Option>, +} + +impl NeverSessionManager { + /// Configure resumable SSE storage without enabling sessions. + pub fn with_event_store(mut self, event_store: Arc) -> Self { + self.event_store = Some(event_store); + self + } +} #[non_exhaustive] pub enum NeverTransport {} impl Transport for NeverTransport { @@ -107,4 +119,8 @@ impl SessionManager for NeverSessionManager { ) -> impl Future> + Send { futures::future::ready(Err(ErrorSessionManagementNotSupported)) } + + fn event_store(&self) -> Option> { + self.event_store.clone() + } } diff --git a/crates/rmcp/src/transport/streamable_http_server/session/store.rs b/crates/rmcp/src/transport/streamable_http_server/session/store.rs index e9a6de2..8b65620 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/store.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/store.rs @@ -1,4 +1,61 @@ -use crate::model::InitializeRequestParams; +use std::pin::Pin; + +use futures::Stream; + +use crate::{ + model::InitializeRequestParams, transport::common::server_side_http::ServerSseMessage, +}; + +/// An opaque identifier for a persisted SSE event. +pub type EventId = String; + +/// An opaque identifier for an SSE stream. +pub type StreamId = String; + +/// A stream of persisted SSE events in delivery order. +pub type EventStream = Pin + Send + Sync + 'static>>; + +/// Type alias for boxed event store errors. +pub type EventStoreError = Box; + +/// Persistent storage for resumable Streamable HTTP events. +/// +/// Implementations typically use a database or distributed log shared by all +/// server instances. Event IDs must be globally unique across all streams, and +/// events must be committed before [`EventStore::store_event`] returns so the +/// returned ID is safe to send to a client. +#[async_trait::async_trait] +pub trait EventStore: Send + Sync + 'static { + /// Persist an event and return the opaque ID clients should receive. + /// + /// The store assigns a globally unique ID and must retain its association + /// with `stream_id` so a later replay only returns events from that stream. + async fn store_event( + &self, + stream_id: &str, + event: &ServerSseMessage, + ) -> Result; + + /// Return events strictly after `last_event_id` in delivery order. + /// + /// Implementations must locate the stream from the globally unique event + /// ID and yield only later events from that stream, with their originally + /// assigned event IDs. + /// + /// A finite stream enables reconnect-and-poll behavior. Implementations + /// backed by a distributed log may keep the stream open to deliver new + /// events as they are appended by any server instance. + async fn replay_events_after( + &self, + last_event_id: &str, + ) -> Result; +} + +impl std::fmt::Debug for dyn EventStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("") + } +} /// State persisted to an external store for cross-instance session recovery. /// diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index 22be737..1b53820 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/crates/rmcp/src/transport/streamable_http_server/tower.rs @@ -1,27 +1,40 @@ use std::{ - borrow::Cow, collections::HashMap, convert::Infallible, fmt::Display, sync::Arc, time::Duration, + borrow::Cow, + collections::HashMap, + convert::Infallible, + fmt::Display, + pin::Pin, + sync::Arc, + task::{Context, Poll}, + time::Duration, }; use bytes::Bytes; -use futures::{StreamExt, future::BoxFuture}; +use futures::{Stream, StreamExt, future::BoxFuture}; use http::{HeaderMap, Method, Request, Response, header::ALLOW}; use http_body::Body; use http_body_util::{BodyExt, Full, combinators::BoxBody}; +use pin_project_lite::pin_project; use tokio_stream::wrappers::ReceiverStream; use tokio_util::sync::CancellationToken; use super::session::{ - RestoreOutcome, SessionId, SessionManager, SessionRestoreMarker, SessionState, SessionStore, + EventStore, EventStoreError, RestoreOutcome, SessionId, SessionManager, SessionRestoreMarker, + SessionState, SessionStore, }; use crate::{ RoleServer, model::{ - ClientCapabilities, ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, - GetExtensions, Implementation, InitializeRequest, InitializeRequestParams, - InitializedNotification, JsonRpcError, ProtocolVersion, RequestId, + ClientCapabilities, ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorCode, + ErrorData, GetExtensions, GetMeta, Implementation, InitializeRequest, + InitializeRequestParams, InitializedNotification, JsonObject, JsonRpcError, + ProtocolVersion, RequestId, ServerConfig, ServerJsonRpcMessage, ServerResult, }, serve_server, - service::serve_directly, + service::{ + NotificationContext, RequestContext, Service, negotiate_protocol_version, + serve_directly_with_ct, uses_legacy_lifecycle, + }, transport::{ OneshotTransport, TransportAdapterIdentity, common::{ @@ -29,6 +42,7 @@ use crate::{ EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_MCP_PROTOCOL_VERSION, HEADER_SESSION_ID, JSON_MIME_TYPE, }, + mcp_headers, server_side_http::{ BoxResponse, ServerSseMessage, accepted_response, expect_json, internal_error_response, sse_stream_response, unexpected_message_response, @@ -37,6 +51,28 @@ use crate::{ }, }; +/// Default maximum POST request body size (4 MiB). +pub(crate) const DEFAULT_MAX_REQUEST_BODY_BYTES: usize = 4 * 1024 * 1024; +const STATELESS_STREAM_CHANNEL_CAPACITY: usize = 16; + +struct ErrorResponse(Box); + +impl ErrorResponse { + fn into_response(self) -> BoxResponse { + *self.0 + } +} + +impl From for ErrorResponse { + fn from(response: BoxResponse) -> Self { + Self(Box::new(response)) + } +} + +type HttpResult = Result; +type RestoreResultSender = tokio::sync::watch::Sender>; +type PendingRestores = Arc>>; + #[non_exhaustive] #[derive(Debug, Clone)] pub struct StreamableHttpServerConfig { @@ -46,11 +82,16 @@ pub struct StreamableHttpServerConfig { pub sse_retry: Option, /// If true, the server will create a session for each request and keep it alive. /// When enabled, SSE priming events are sent to enable client reconnection. - pub stateful_mode: bool, - /// When true and `stateful_mode` is false, the server returns - /// `Content-Type: application/json` directly instead of `text/event-stream`. - /// This eliminates SSE framing overhead for simple request-response tools, - /// allowed by the MCP Streamable HTTP spec (2025-06-18). + /// + /// Only applies to legacy protocol versions (`< 2026-07-28`). Per SEP-2567, + /// sessions are removed from the `2026-07-28` version, so requests + /// negotiating that version are always served statelessly regardless of + /// this setting. + pub legacy_session_mode: bool, + /// When true and `legacy_session_mode` is false, the server prefers + /// `Content-Type: application/json` for simple request-response tools. + /// If the handler emits a notification or request before the final response, + /// the server falls back to `text/event-stream` so no message is lost. pub json_response: bool, /// Cancellation token for the Streamable HTTP server. /// @@ -69,13 +110,18 @@ pub struct StreamableHttpServerConfig { pub allowed_hosts: Vec, /// Allowed browser origins for inbound `Origin` validation. /// - /// Defaults to an empty list, which disables Origin validation. When - /// non-empty, requests carrying an `Origin` header must match per RFC 6454 - /// `(scheme, host, port)`; missing-`Origin` requests still pass. Entries - /// must include a scheme; `"null"` matches the browser's `Origin: null`. + /// Defaults to an empty list, which disables Origin validation for backward + /// compatibility. A non-empty list enables validation. Requests carrying + /// an `Origin` header must match per RFC 6454 `(scheme, host, port)`; + /// missing-`Origin` requests still pass. Entries must include a scheme; + /// `"null"` matches the browser's `Origin: null`. + /// + /// Call [`StreamableHttpServerConfig::enforce_origin_validation`] to enable + /// validation with an empty list, rejecting every present Origin value. /// examples: /// allowed_origins = ["https://app.example.com", "http://localhost:8080"] pub allowed_origins: Vec, + validate_empty_origin_allowlist: bool, /// Optional external session store for cross-instance recovery. /// /// When set, [`SessionState`] (the client's `initialize` parameters) is @@ -98,6 +144,38 @@ pub struct StreamableHttpServerConfig { /// }; /// ``` pub session_store: Option>, + /// Maximum POST request body size in bytes. + /// + /// Enforced while streaming the body, independent of `Content-Length`, + /// chunked transfer encoding, or HTTP version. Oversized payloads receive + /// a `413 Payload Too Large` response. + pub max_request_body_bytes: usize, + /// Require stateless JSON-RPC request POSTs to carry per-request protocol + /// signals before handler dispatch. + /// + /// Non-initialize requests must carry `MCP-Protocol-Version`; ordinary + /// non-discovery requests must also carry + /// `_meta.io.modelcontextprotocol/protocolVersion`. `server/discover` + /// retains its existing request-metadata validation. For `2026-07-28` + /// requests, the server handler continues to require the remaining + /// per-request metadata, including `clientCapabilities`. Initialize, + /// notifications, and other message kinds retain their existing rules. + /// + /// This option applies to requests routed statelessly. Set + /// `legacy_session_mode` to `false` to ensure every request uses that path. + /// Legacy session routing and its error precedence remain unchanged. + /// + /// The validator checks metadata presence rather than applying a version + /// allowlist. However, rmcp clients negotiated below `2026-07-28` do not + /// attach per-request protocol metadata, so enabling this option rejects + /// their ordinary requests. Servers using this option should normally + /// override + /// [`ServerHandler::supported_protocol_versions`](crate::ServerHandler::supported_protocol_versions) + /// to advertise only `2026-07-28` and later. + /// + /// Default is `false`, preserving today's legacy behavior where an absent + /// header is treated as protocol version `2025-03-26`. + pub stateless_protocol_metadata_required: bool, } impl std::fmt::Debug for dyn SessionStore { @@ -111,12 +189,15 @@ impl Default for StreamableHttpServerConfig { Self { sse_keep_alive: Some(Duration::from_secs(15)), sse_retry: Some(Duration::from_secs(3)), - stateful_mode: true, + legacy_session_mode: true, json_response: false, cancellation_token: CancellationToken::new(), allowed_hosts: vec!["localhost".into(), "127.0.0.1".into(), "::1".into()], allowed_origins: vec![], + validate_empty_origin_allowlist: false, session_store: None, + max_request_body_bytes: DEFAULT_MAX_REQUEST_BODY_BYTES, + stateless_protocol_metadata_required: false, } } } @@ -141,9 +222,15 @@ impl StreamableHttpServerConfig { self.allowed_origins = allowed_origins.into_iter().map(Into::into).collect(); self } - /// Disable Origin validation, reverting to the default ignore-Origin behavior. + /// Enable Origin validation, including when the allowed Origins list is empty. + pub fn enforce_origin_validation(mut self) -> Self { + self.validate_empty_origin_allowlist = true; + self + } + /// Disable Origin validation, allowing requests with any `Origin` header. pub fn disable_allowed_origins(mut self) -> Self { self.allowed_origins.clear(); + self.validate_empty_origin_allowlist = false; self } pub fn with_sse_keep_alive(mut self, duration: Option) -> Self { @@ -156,8 +243,8 @@ impl StreamableHttpServerConfig { self } - pub fn with_stateful_mode(mut self, stateful: bool) -> Self { - self.stateful_mode = stateful; + pub fn with_legacy_session_mode(mut self, legacy_session_mode: bool) -> Self { + self.legacy_session_mode = legacy_session_mode; self } @@ -170,18 +257,35 @@ impl StreamableHttpServerConfig { self.cancellation_token = token; self } + + /// Set the maximum POST request body size in bytes. + pub fn with_max_request_body_bytes(mut self, bytes: usize) -> Self { + self.max_request_body_bytes = bytes; + self + } + + /// Require per-request protocol signals on stateless JSON-RPC request + /// POSTs. + /// + /// See [`StreamableHttpServerConfig::stateless_protocol_metadata_required`]. + pub fn with_stateless_protocol_metadata_required( + mut self, + stateless_protocol_metadata_required: bool, + ) -> Self { + self.stateless_protocol_metadata_required = stateless_protocol_metadata_required; + self + } } -#[expect( - clippy::result_large_err, - reason = "BoxResponse is intentionally large; matches other handlers in this file" -)] /// Validates the `MCP-Protocol-Version` header on incoming HTTP requests. /// /// Per the MCP 2025-06-18 spec: /// - If the header is present but contains an unsupported version, return 400 Bad Request. /// - If the header is absent, assume `2025-03-26` for backwards compatibility (no error). -fn validate_protocol_version_header(headers: &http::HeaderMap) -> Result<(), BoxResponse> { +fn validate_protocol_version_header( + headers: &http::HeaderMap, + allow_unknown: bool, +) -> HttpResult<()> { if let Some(value) = headers.get(HEADER_MCP_PROTOCOL_VERSION) { let version_str = value.to_str().map_err(|_| { Response::builder() @@ -197,7 +301,7 @@ fn validate_protocol_version_header(headers: &http::HeaderMap) -> Result<(), Box let is_known = ProtocolVersion::KNOWN_VERSIONS .iter() .any(|v| v.as_str() == version_str); - if !is_known { + if !allow_unknown && !is_known { return Err(Response::builder() .status(http::StatusCode::BAD_REQUEST) .body( @@ -206,12 +310,157 @@ fn validate_protocol_version_header(headers: &http::HeaderMap) -> Result<(), Box ))) .boxed(), ) - .expect("valid response")); + .expect("valid response") + .into()); } } Ok(()) } +fn message_has_per_request_protocol_version(message: &ClientJsonRpcMessage) -> bool { + match message { + ClientJsonRpcMessage::Request(request) => { + request.request.get_meta().protocol_version().is_some() + } + _ => false, + } +} + +struct NegotiatingStatelessHttpService(S); + +impl> Service for NegotiatingStatelessHttpService { + async fn handle_request( + &self, + request: ClientRequest, + context: RequestContext, + ) -> Result { + let requested_protocol_version = + if let ClientRequest::InitializeRequest(initialize) = &request { + Some(initialize.params.protocol_version.clone()) + } else { + None + }; + let peer = context.peer.clone(); + let mut response = self.0.handle_request(request, context).await?; + if let (Some(requested), ServerResult::InitializeResult(result)) = + (requested_protocol_version, &mut response) + { + result.protocol_version = negotiate_protocol_version( + &requested, + result.protocol_version.clone(), + &self.0.supported_protocol_versions(), + )?; + if let Some(peer_info) = peer.peer_info() { + let mut peer_info = (*peer_info).clone(); + peer_info.protocol_version = result.protocol_version.clone(); + peer.set_peer_info(peer_info); + } + } + Ok(response) + } + + async fn handle_notification( + &self, + notification: ClientNotification, + context: NotificationContext, + ) -> Result<(), ErrorData> { + self.0.handle_notification(notification, context).await + } + + fn get_info(&self) -> ServerConfig { + self.0.get_info() + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + self.0.supported_protocol_versions() + } +} + +// SEP-2567: sessions are removed from the discover lifecycle. Validate +// protocol-version consistency, then classify the request with the shared +// lifecycle helper. +fn is_legacy_request( + message: Option<&ClientJsonRpcMessage>, + headers: &HeaderMap, +) -> HttpResult { + let has_per_request_version = message.is_some_and(message_has_per_request_protocol_version); + validate_protocol_version_header(headers, has_per_request_version)?; + if let Some(message) = message { + if let ClientJsonRpcMessage::Request(req) = message + && let ClientRequest::InitializeRequest(init) = &req.request + { + validate_header_matches_init_body( + headers, + init.params.protocol_version.as_str(), + Some(req.id.clone()), + )?; + } + validate_request_protocol_version_meta(headers, message)?; + } + + // An `initialize` request selects legacy semantics whatever version it names: + // the handshake exists only in the revisions before 2026-07-28, so the + // version in its params never routes it to the stateless path. The + // handshake itself answers with a legacy version the server supports. + if matches!( + message, + Some(ClientJsonRpcMessage::Request(req)) + if matches!(&req.request, ClientRequest::InitializeRequest(_)) + ) { + return Ok(true); + } + + let uses_discover_lifecycle = matches!( + message, + Some(ClientJsonRpcMessage::Request(req)) + if req + .request + .get_meta() + .missing_required_keys(&ProtocolVersion::V_2026_07_28) + .is_empty() + ); + + let from_body = match message { + Some(ClientJsonRpcMessage::Request(req)) => req.request.get_meta().protocol_version(), + _ => None, + }; + let version = from_body + .or_else(|| { + headers + .get(HEADER_MCP_PROTOCOL_VERSION) + .and_then(|value| value.to_str().ok()) + .and_then(|s| serde_json::from_value(serde_json::Value::String(s.to_owned())).ok()) + }) + .unwrap_or(ProtocolVersion::V_2025_03_26); + Ok(uses_legacy_lifecycle( + Some(&version), + uses_discover_lifecycle, + )) +} + +fn method_not_allowed_response() -> BoxResponse { + Response::builder() + .status(http::StatusCode::METHOD_NOT_ALLOWED) + .header(ALLOW, "POST") + .body(Full::new(Bytes::from("Method Not Allowed")).boxed()) + .expect("valid response") +} + +async fn persist_and_forward_event( + event_store: &dyn EventStore, + stream_id: &str, + mut event: ServerSseMessage, + output: &mut Option>, +) -> Result<(), EventStoreError> { + event.event_id = Some(event_store.store_event(stream_id, &event).await?); + if let Some(sender) = output + && sender.send(event).await.is_err() + { + *output = None; + } + Ok(()) +} + fn invalid_request_jsonrpc_response( id: Option, message: impl Into>, @@ -225,16 +474,25 @@ fn invalid_request_jsonrpc_response( .expect("valid response") } -#[expect( - clippy::result_large_err, - reason = "BoxResponse is intentionally large; matches other handlers in this file" -)] +fn invalid_params_jsonrpc_response( + id: Option, + message: impl Into>, +) -> BoxResponse { + let err = JsonRpcError::new(id, ErrorData::invalid_params(message, None)); + let body = serde_json::to_vec(&err).expect("serialize JsonRpcError"); + Response::builder() + .status(http::StatusCode::BAD_REQUEST) + .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE) + .body(Full::new(Bytes::from(body)).boxed()) + .expect("valid response") +} + /// Absent header is allowed; the first initialize round-trip may legitimately omit it. fn validate_header_matches_init_body( headers: &http::HeaderMap, body_version: &str, request_id: Option, -) -> Result<(), BoxResponse> { +) -> HttpResult<()> { let Some(header_value) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else { return Ok(()); }; @@ -255,7 +513,258 @@ fn validate_header_matches_init_body( format!( "Invalid Request: MCP-Protocol-Version header ({header_str}) does not match initialize params.protocolVersion ({body_version})" ), - )); + ) + .into()); + } + Ok(()) +} + +fn validate_request_protocol_version_meta( + headers: &HeaderMap, + message: &ClientJsonRpcMessage, +) -> HttpResult<()> { + let ClientJsonRpcMessage::Request(request) = message else { + return Ok(()); + }; + if matches!(&request.request, ClientRequest::InitializeRequest(_)) { + return Ok(()); + } + let is_discover = matches!(&request.request, ClientRequest::DiscoverRequest(_)); + let meta = request.request.get_meta(); + let header_version = headers + .get(HEADER_MCP_PROTOCOL_VERSION) + .and_then(|value| value.to_str().ok()); + let Some(meta_version) = meta.protocol_version() else { + let requires_request_metadata = is_discover + || header_version + .is_some_and(|version| version >= ProtocolVersion::V_2026_07_28.as_str()); + if requires_request_metadata { + let missing = meta.missing_required_keys(&ProtocolVersion::V_2026_07_28); + return Err(invalid_params_jsonrpc_response( + Some(request.id.clone()), + format!( + "Invalid params: request _meta is missing or has malformed required fields: {}", + missing.join(", ") + ), + ) + .into()); + } + return Ok(()); + }; + let Some(header_version) = header_version else { + return Err(header_mismatch_jsonrpc_response( + Some(request.id.clone()), + "request _meta protocolVersion requires MCP-Protocol-Version header", + ) + .into()); + }; + if header_version != meta_version.as_str() { + return Err(header_mismatch_jsonrpc_response( + Some(request.id.clone()), + format!( + "MCP-Protocol-Version header ({header_version}) does not match request _meta protocolVersion ({meta_version})" + ), + ) + .into()); + } + Ok(()) +} + +/// When `stateless_protocol_metadata_required` is enabled in stateless mode, +/// every non-initialize Streamable HTTP JSON-RPC request POST must carry the +/// `MCP-Protocol-Version` HTTP header. A missing header is rejected with +/// HTTP 400 / JSON-RPC `-32020` before handler dispatch. `server/discover` +/// is included so the seam aligns with the per-POST header contract; its +/// body-metadata rule is preserved unchanged. +fn validate_required_protocol_header( + config: &StreamableHttpServerConfig, + headers: &HeaderMap, + message: &ClientJsonRpcMessage, +) -> HttpResult<()> { + if !config.stateless_protocol_metadata_required { + return Ok(()); + } + let ClientJsonRpcMessage::Request(request) = message else { + // Notifications, response messages, and error messages are exempt. + return Ok(()); + }; + if matches!(&request.request, ClientRequest::InitializeRequest(_)) { + // Initialize keeps its own header-matching rule. + return Ok(()); + } + if headers.contains_key(HEADER_MCP_PROTOCOL_VERSION) { + return Ok(()); + } + Err(header_mismatch_jsonrpc_response( + Some(request.id.clone()), + "Missing MCP-Protocol-Version header for request requiring per-request protocol metadata", + ) + .into()) +} + +/// When `stateless_protocol_metadata_required` is enabled in stateless mode, +/// every non-initialize, non-discover Streamable HTTP JSON-RPC request must +/// carry `io.modelcontextprotocol/protocolVersion` in `_meta`. A missing entry +/// is rejected with HTTP 400 / JSON-RPC `-32602` (invalid_params). `initialize`, +/// `server/discover` (whose body-metadata rule is already enforced by +/// `validate_request_protocol_version_meta`), notifications, and other message +/// kinds are exempt. +fn validate_required_protocol_meta( + config: &StreamableHttpServerConfig, + message: &ClientJsonRpcMessage, +) -> HttpResult<()> { + if !config.stateless_protocol_metadata_required { + return Ok(()); + } + let ClientJsonRpcMessage::Request(request) = message else { + return Ok(()); + }; + if matches!( + &request.request, + ClientRequest::InitializeRequest(_) | ClientRequest::DiscoverRequest(_) + ) { + return Ok(()); + } + if request.request.get_meta().protocol_version().is_some() { + return Ok(()); + } + Err(invalid_params_jsonrpc_response( + Some(request.id.clone()), + "Invalid params: request requires protocolVersion in request _meta", + ) + .into()) +} + +fn jsonrpc_http_status(message: &ServerJsonRpcMessage) -> http::StatusCode { + let ServerJsonRpcMessage::Error(error) = message else { + return http::StatusCode::OK; + }; + // Modern per-request HTTP treats invalid params as a malformed request. + // Legacy requests bypass this mapper and retain HTTP 200 JSON-RPC errors. + match error.error.code { + ErrorCode::UNSUPPORTED_PROTOCOL_VERSION + | ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY + | ErrorCode::INVALID_PARAMS + | ErrorCode::HEADER_MISMATCH => http::StatusCode::BAD_REQUEST, + ErrorCode::METHOD_NOT_FOUND => http::StatusCode::NOT_FOUND, + _ => http::StatusCode::OK, + } +} + +#[cfg(test)] +mod jsonrpc_http_status_tests { + use super::*; + + fn error_message(code: ErrorCode) -> ServerJsonRpcMessage { + ServerJsonRpcMessage::Error(JsonRpcError::new( + Some(RequestId::Number(1)), + ErrorData::new(code, "test error", None), + )) + } + + /// A handler returning `ErrorData::header_mismatch(..)` on the modern per-request path + /// must map to HTTP 400, not the default HTTP 200. Regression test for + /// https://github.com/modelcontextprotocol/rust-sdk/issues/1225 + #[test] + fn header_mismatch_maps_to_bad_request() { + assert_eq!( + jsonrpc_http_status(&error_message(ErrorCode::HEADER_MISMATCH)), + http::StatusCode::BAD_REQUEST + ); + } + + #[test] + fn method_not_found_maps_to_not_found() { + assert_eq!( + jsonrpc_http_status(&error_message(ErrorCode::METHOD_NOT_FOUND)), + http::StatusCode::NOT_FOUND + ); + } + + #[test] + fn unmapped_error_defaults_to_ok() { + assert_eq!( + jsonrpc_http_status(&error_message(ErrorCode::INTERNAL_ERROR)), + http::StatusCode::OK + ); + } +} + +fn jsonrpc_message_response( + message: ServerJsonRpcMessage, + map_protocol_status: bool, +) -> HttpResult { + let status = if map_protocol_status { + jsonrpc_http_status(&message) + } else { + http::StatusCode::OK + }; + let body = + serde_json::to_vec(&message).map_err(internal_error_response("serialize json response"))?; + Ok(Response::builder() + .status(status) + .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE) + .body(Full::new(Bytes::from(body)).boxed()) + .expect("valid response")) +} + +fn header_mismatch_jsonrpc_response( + id: Option, + message: impl Into>, +) -> BoxResponse { + let err = JsonRpcError::new(id, ErrorData::header_mismatch(message, None)); + let body = serde_json::to_vec(&err).expect("serialize JsonRpcError"); + Response::builder() + .status(http::StatusCode::BAD_REQUEST) + .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE) + .body(Full::new(Bytes::from(body)).boxed()) + .expect("valid response") +} + +/// Validates SEP-2243 `Mcp-Method` / `Mcp-Name` / `Mcp-Param-*` headers against the body. +/// +/// Only enforced when the request declares a protocol version `>= STANDARD_HEADERS`. +/// The `initialize` handshake is exempt: clients emit these headers only after the +/// version has been negotiated. `tool_schema` supplies the called tool's input schema +/// so annotated `Mcp-Param-*` headers can be checked (no schema => those are skipped). +fn validate_standard_headers( + headers: &HeaderMap, + message: &ClientJsonRpcMessage, + tool_schema: impl Fn(&str) -> Option>, +) -> HttpResult<()> { + let version_requires_headers = headers + .get(HEADER_MCP_PROTOCOL_VERSION) + .and_then(|value| value.to_str().ok()) + .is_some_and(|version| version >= ProtocolVersion::STANDARD_HEADERS.as_str()); + if !version_requires_headers { + return Ok(()); + } + + let request_id = match message { + ClientJsonRpcMessage::Request(req) => { + if matches!(&req.request, ClientRequest::InitializeRequest(_)) { + return Ok(()); + } + Some(req.id.clone()) + } + ClientJsonRpcMessage::Notification(_) => None, + _ => return Ok(()), + }; + + let Ok(value) = serde_json::to_value(message) else { + return Ok(()); + }; + // For tools/call, look up the tool schema so Mcp-Param-* headers are validated. + let schema = value + .get("method") + .and_then(|method| method.as_str()) + .filter(|method| *method == "tools/call") + .and_then(|_| value.get("params")) + .and_then(|params| params.get("name")) + .and_then(|name| name.as_str()) + .and_then(tool_schema); + if let Err(reason) = mcp_headers::validate_request_headers(headers, &value, schema.as_deref()) { + return Err(header_mismatch_jsonrpc_response(request_id, reason).into()); } Ok(()) } @@ -345,9 +854,6 @@ fn parse_origin_value(value: &str) -> Option { } fn origin_is_allowed(origin: &NormalizedOrigin, allowed_origins: &[String]) -> bool { - if allowed_origins.is_empty() { - return true; - } allowed_origins .iter() .filter_map(|raw| parse_origin_value(raw)) @@ -379,10 +885,7 @@ fn bad_request_response(message: &str) -> BoxResponse { .expect("failed to build bad request response") } -fn parse_host_header( - uri: &http::Uri, - headers: &HeaderMap, -) -> Result { +fn parse_host_header(uri: &http::Uri, headers: &HeaderMap) -> HttpResult { if let Some(host) = headers.get(http::header::HOST) { let host_str = host .to_str() @@ -413,24 +916,24 @@ fn validate_dns_rebinding_headers( uri: &http::Uri, headers: &HeaderMap, config: &StreamableHttpServerConfig, -) -> Result<(), BoxResponse> { +) -> HttpResult<()> { let host = parse_host_header(uri, headers)?; if !host_is_allowed(&host, &config.allowed_hosts) { tracing::warn!( host = ?host, "rejected request with disallowed Host header (possible DNS rebinding attempt)", ); - return Err(forbidden_response("Forbidden: Host header is not allowed")); + return Err(forbidden_response("Forbidden: Host header is not allowed").into()); } - validate_origin_header(headers, &config.allowed_origins)?; + validate_origin_header(headers, config)?; Ok(()) } fn validate_origin_header( headers: &HeaderMap, - allowed_origins: &[String], -) -> Result<(), BoxResponse> { - if allowed_origins.is_empty() { + config: &StreamableHttpServerConfig, +) -> HttpResult<()> { + if !config.validate_empty_origin_allowlist && config.allowed_origins.is_empty() { return Ok(()); } let Some(origin_header) = headers.get(http::header::ORIGIN) else { @@ -441,22 +944,20 @@ fn validate_origin_header( .inspect_err(|_| { tracing::warn!(origin = ?origin_header, "rejected request with non-UTF-8 Origin header"); }) - .map_err(|_| bad_request_response("Bad Request: Invalid Origin header encoding"))?; + .map_err(|_| forbidden_response("Forbidden: Invalid Origin header encoding"))?; let origin = parse_origin_value(origin_str).ok_or_else(|| { tracing::warn!( origin = origin_str, "rejected request with malformed Origin header", ); - bad_request_response("Bad Request: Invalid Origin header") + forbidden_response("Forbidden: Invalid Origin header") })?; - if !origin_is_allowed(&origin, allowed_origins) { + if !origin_is_allowed(&origin, &config.allowed_origins) { tracing::warn!( origin = ?origin, "rejected request with disallowed Origin header (possible cross-origin attack)", ); - return Err(forbidden_response( - "Forbidden: Origin header is not allowed", - )); + return Err(forbidden_response("Forbidden: Origin header is not allowed").into()); } Ok(()) } @@ -464,12 +965,12 @@ fn validate_origin_header( /// # Streamable HTTP server /// /// An HTTP service that implements the -/// [Streamable HTTP transport](https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#streamable-http) +/// [Streamable HTTP transport](https://modelcontextprotocol.io/specification/2026-07-28/basic/transports/streamable-http) /// for MCP servers. /// /// ## Session management /// -/// When [`StreamableHttpServerConfig::stateful_mode`] is `true` (the default), +/// When [`StreamableHttpServerConfig::legacy_session_mode`] is `true` (the default), /// the server creates a session for each client that sends an `initialize` /// request. The session ID is returned in the `Mcp-Session-Id` response header /// and the client must include it on all subsequent requests. @@ -552,9 +1053,11 @@ pub struct StreamableHttpService { /// same unknown session ID wait for the first restore to complete rather /// than racing to replay the initialize handshake. `None` when no external /// session store is configured (avoids allocating the map). - pending_restores: Option< - Arc>>>>, - >, + pending_restores: Option, + /// Caches tool input schemas by name for SEP-2243 `Mcp-Param-*` validation. + /// Populated lazily via `get_tool` so the service factory runs at most once + /// per tool name. `None` value means the tool exposes no schema. + tool_schemas: Arc>>>>, } impl Clone for StreamableHttpService { @@ -564,6 +1067,7 @@ impl Clone for StreamableHttpService { session_manager: self.session_manager.clone(), service_factory: self.service_factory.clone(), pending_restores: self.pending_restores.clone(), + tool_schemas: self.tool_schemas.clone(), } } } @@ -571,7 +1075,7 @@ impl Clone for StreamableHttpService { impl tower_service::Service> for StreamableHttpService where RequestBody: Body + Send + 'static, - S: crate::Service + Send + 'static, + S: crate::ServerHandler + Send + 'static, M: SessionManager, RequestBody::Error: Display, RequestBody::Data: Send + 'static, @@ -602,10 +1106,9 @@ where /// `result` defaults to `false` (failure / cancellation). Only the success path /// needs to set it to `true` before returning. struct PendingRestoreGuard { - pending_restores: - Arc>>>>, + pending_restores: PendingRestores, session_id: SessionId, - watch_tx: tokio::sync::watch::Sender>, + watch_tx: RestoreResultSender, /// The value that will be broadcast to waiting tasks on drop. result: bool, } @@ -625,7 +1128,7 @@ impl Drop for PendingRestoreGuard { impl StreamableHttpService where - S: crate::Service + Send + 'static, + S: crate::ServerHandler + Send + 'static, M: SessionManager, { pub fn new( @@ -633,23 +1136,200 @@ where session_manager: Arc, config: StreamableHttpServerConfig, ) -> Self { - let pending_restores = config.session_store.is_some().then(|| { - Arc::new(tokio::sync::RwLock::new(HashMap::< - SessionId, - tokio::sync::watch::Sender>, - >::new())) - }); + let pending_restores = config + .session_store + .is_some() + .then(|| Arc::new(tokio::sync::RwLock::new(HashMap::new()))); Self { config, session_manager, service_factory: Arc::new(service_factory), pending_restores, + tool_schemas: Arc::new(std::sync::RwLock::new(HashMap::new())), } } fn get_service(&self) -> Result { (self.service_factory)() } + fn persisted_stateless_stream( + &self, + first: Option, + mut receiver: tokio::sync::mpsc::Receiver, + request_ct: CancellationToken, + event_store: Arc, + ) -> ReceiverStream { + let (sender, output) = tokio::sync::mpsc::channel(STATELESS_STREAM_CHANNEL_CAPACITY); + let stream_id = uuid::Uuid::new_v4().to_string(); + let retry = self.config.sse_retry; + let server_ct = self.config.cancellation_token.child_token(); + + tokio::spawn(async move { + let mut sender = Some(sender); + if let Some(retry) = retry + && let Err(error) = persist_and_forward_event( + event_store.as_ref(), + &stream_id, + ServerSseMessage::retry(retry), + &mut sender, + ) + .await + { + tracing::error!(%stream_id, %error, "failed to persist SSE priming event"); + request_ct.cancel(); + return; + } + + let mut first = first; + loop { + let message = if let Some(message) = first.take() { + Some(message) + } else { + tokio::select! { + message = receiver.recv() => message, + _ = server_ct.cancelled() => { + request_ct.cancel(); + None + } + } + }; + let Some(message) = message else { + break; + }; + tracing::trace!(?message); + if let Err(error) = persist_and_forward_event( + event_store.as_ref(), + &stream_id, + ServerSseMessage::from_message(message), + &mut sender, + ) + .await + { + tracing::error!(%stream_id, %error, "failed to persist SSE event"); + request_ct.cancel(); + break; + } + } + }); + + ReceiverStream::new(output) + } + + fn stateless_sse_response( + &self, + first: Option, + receiver: tokio::sync::mpsc::Receiver, + request_ct: CancellationToken, + ) -> BoxResponse { + if let Some(event_store) = self.session_manager.event_store() { + let stream = self.persisted_stateless_stream(first, receiver, request_ct, event_store); + sse_stream_response( + stream, + self.config.sse_keep_alive, + self.config.cancellation_token.child_token(), + ) + } else { + let stream = futures::stream::iter(first) + .chain(ReceiverStream::new(receiver)) + .map(|message| { + tracing::trace!(?message); + ServerSseMessage::from_message(message) + }); + sse_stream_response( + CancelOnDisconnect::new(stream, request_ct), + self.config.sse_keep_alive, + self.config.cancellation_token.child_token(), + ) + } + } + + // The HTTP status must be known before opening an SSE stream. + async fn serve_negotiated_request_directly( + &self, + service: S, + mut request: crate::model::JsonRpcRequest, + parts: http::request::Parts, + ) -> HttpResult { + let peer_info = Self::peer_info_for_stateless_request(&request, &parts.headers); + request.request.extensions_mut().insert(parts); + let (transport, mut receiver) = + OneshotTransport::::new(ClientJsonRpcMessage::Request(request)); + // Give this stateless request its own cancellation token so a client + // disconnect can cancel the in-flight handler (#857), as in the + // non-negotiated stateless path below. + let request_ct = CancellationToken::new(); + let service = serve_directly_with_ct( + NegotiatingStatelessHttpService(service), + transport, + peer_info, + request_ct.clone(), + ); + tokio::spawn(async move { + let _ = service.waiting().await; + }); + + let cancel = self.config.cancellation_token.child_token(); + // Cancel the handler if the client disconnects while it is still + // producing its first message (this future is dropped before + // `receiver.recv()` completes). Disarmed once the handler emits + // anything, so a normal response is never cancelled. + let mut disconnect_guard = Some(request_ct.clone().drop_guard()); + let first = tokio::select! { + message = receiver.recv() => { + if let Some(guard) = disconnect_guard.take() { + guard.disarm(); + } + message + } + _ = cancel.cancelled() => None, + } + .ok_or_else(|| { + internal_error_response("empty response")(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "no response message received from handler", + )) + })?; + + let terminal = matches!( + &first, + ServerJsonRpcMessage::Response(_) | ServerJsonRpcMessage::Error(_) + ); + if terminal + && (self.config.json_response || jsonrpc_http_status(&first) != http::StatusCode::OK) + { + // This message is the whole reply, so `receiver` is dropped here and + // anything the handler emits afterwards is undeliverable. Cancel it so + // a still-running handler stops instead of running on unobserved: its + // terminal `send` would otherwise fail before adding the termination + // permit, leaving the serve loop parked forever. A no-op when the + // handler already completed. + request_ct.cancel(); + return jsonrpc_message_response(first, true); + } + + Ok(self.stateless_sse_response(Some(first), receiver, request_ct)) + } + + /// Returns the cached input schema for `name`, constructing a service once + /// per name to read its `ServerHandler::get_tool` definition. Used to + /// validate SEP-2243 `Mcp-Param-*` headers against the request body. + fn tool_schema(&self, name: &str) -> Option> { + if let Ok(cache) = self.tool_schemas.read() + && let Some(schema) = cache.get(name) + { + return schema.clone(); + } + let schema = self + .get_service() + .ok() + .and_then(|service| service.get_tool(name)) + .map(|tool| tool.input_schema); + if let Ok(mut cache) = self.tool_schemas.write() { + cache.insert(name.to_owned(), schema.clone()); + } + schema + } + /// Spawn a task that runs `serve_server` for the given session, waits for /// it to finish, and then calls `close_session`. /// @@ -664,7 +1344,7 @@ where transport: M::Transport, init_done_tx: Option>, ) where - S: crate::Service + Send + 'static, + S: crate::ServerHandler + Send + 'static, M: SessionManager, { tokio::spawn(async move { @@ -707,7 +1387,7 @@ where parts: &http::request::Parts, ) -> Result where - S: crate::Service + Send + 'static, + S: crate::ServerHandler + Send + 'static, M: SessionManager, { // Both fields are Some iff a session store is configured. @@ -827,23 +1507,15 @@ where Some(init_done_tx), ); - if let Err(e) = self - .session_manager + self.session_manager .initialize_session(session_id, restore_init) .await - .map_err(|e| std::io::Error::other(e.to_string())) - { - return Err(e); - } + .map_err(|e| std::io::Error::other(e.to_string()))?; - if let Err(e) = self - .session_manager + self.session_manager .accept_message(session_id, restore_initialized) .await - .map_err(|e| std::io::Error::other(e.to_string())) - { - return Err(e); - } + .map_err(|e| std::io::Error::other(e.to_string()))?; if init_done_rx.await.is_err() { return Err(std::io::Error::other( @@ -868,18 +1540,21 @@ where if let Err(response) = validate_dns_rebinding_headers(request.uri(), request.headers(), &self.config) { - return response; + return response.into_response(); } let method = request.method().clone(); - let allowed_methods = match self.config.stateful_mode { - true => "GET, POST, DELETE", - false => "POST", + let supports_stateless_replay = self.session_manager.event_store().is_some(); + let allowed_methods = match (self.config.legacy_session_mode, supports_stateless_replay) { + (true, _) => "GET, POST, DELETE", + (false, true) => "GET, POST", + (false, false) => "POST", }; - let result = match (method, self.config.stateful_mode) { - (Method::POST, _) => self.handle_post(request).await, - // if we're not in stateful mode, we don't support GET or DELETE because there is no session - (Method::GET, true) => self.handle_get(request).await, - (Method::DELETE, true) => self.handle_delete(request).await, + let result = match method { + Method::POST => self.handle_post(request).await, + Method::GET if self.config.legacy_session_mode || supports_stateless_replay => { + self.handle_get(request).await + } + Method::DELETE if self.config.legacy_session_mode => self.handle_delete(request).await, _ => { // Handle other methods or return an error let response = Response::builder() @@ -892,10 +1567,10 @@ where }; match result { Ok(response) => response, - Err(response) => response, + Err(response) => response.into_response(), } } - async fn handle_get(&self, request: Request) -> Result + async fn handle_get(&self, request: Request) -> HttpResult where B: Body + Send + 'static, B::Error: Display, @@ -917,6 +1592,32 @@ where ) .expect("valid response")); } + let request_uses_legacy_protocol = is_legacy_request(None, request.headers())?; + let legacy_request = self.config.legacy_session_mode && request_uses_legacy_protocol; + if !legacy_request { + let Some(last_event_id) = request + .headers() + .get(HEADER_LAST_EVENT_ID) + .and_then(|value| value.to_str().ok()) + else { + return Ok(method_not_allowed_response()); + }; + let Some(event_store) = self.session_manager.event_store() else { + return Ok(method_not_allowed_response()); + }; + let stream = match event_store.replay_events_after(last_event_id).await { + Ok(stream) => stream, + Err(error) => { + tracing::warn!(%error, "stateless SSE resume failed, returning empty stream"); + Box::pin(futures::stream::empty()) + } + }; + return Ok(sse_stream_response( + stream, + self.config.sse_keep_alive, + self.config.cancellation_token.child_token(), + )); + } // check session id let session_id = request .headers() @@ -952,7 +1653,7 @@ where } } // Validate MCP-Protocol-Version header (per 2025-06-18 spec) - validate_protocol_version_header(&parts.headers)?; + validate_protocol_version_header(&parts.headers, false)?; // check if last event id is provided let last_event_id = parts .headers @@ -994,7 +1695,11 @@ where .await .map_err(internal_error_response("create standalone stream"))?; let stream = if let Some(retry) = self.config.sse_retry { - let priming = ServerSseMessage::priming("0", retry); + let priming = if self.session_manager.event_store().is_some() { + ServerSseMessage::retry(retry) + } else { + ServerSseMessage::priming("0", retry) + }; futures::stream::once(async move { priming }) .chain(stream) .left_stream() @@ -1008,7 +1713,7 @@ where )) } - async fn handle_post(&self, request: Request) -> Result + async fn handle_post(&self, request: Request) -> HttpResult where B: Body + Send + 'static, B::Error: Display, @@ -1048,12 +1753,15 @@ where // json deserialize request body let (part, body) = request.into_parts(); - let mut message = match expect_json(body).await { + let mut message = match expect_json(body, self.config.max_request_body_bytes).await { Ok(message) => message, Err(response) => return Ok(response), }; - if self.config.stateful_mode { + let use_session = + self.config.legacy_session_mode && is_legacy_request(Some(&message), &part.headers)?; + + if use_session { // do we have a session id? let session_id = part .headers @@ -1082,7 +1790,11 @@ where } // Validate MCP-Protocol-Version header (per 2025-06-18 spec) - validate_protocol_version_header(&part.headers)?; + let has_per_request_version = message_has_per_request_protocol_version(&message); + validate_protocol_version_header(&part.headers, has_per_request_version)?; + validate_request_protocol_version_meta(&part.headers, &message)?; + // Validate SEP-2243 standard headers against the body + validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?; // inject request part to extensions match &mut message { @@ -1125,12 +1837,35 @@ where } } } else { + if matches!( + &message, + ClientJsonRpcMessage::Request(request) + if matches!(&request.request, ClientRequest::DiscoverRequest(_)) + ) { + validate_protocol_version_header( + &part.headers, + message_has_per_request_protocol_version(&message), + )?; + validate_standard_headers(&part.headers, &message, |name| { + self.tool_schema(name) + })?; + validate_request_protocol_version_meta(&part.headers, &message)?; + let ClientJsonRpcMessage::Request(request) = message else { + unreachable!("guarded as a request above"); + }; + let service = self + .get_service() + .map_err(internal_error_response("get service"))?; + return self + .serve_negotiated_request_directly(service, request, part) + .await; + } // Capture init params for external store persistence before // extensions are injected (which would require Clone). let stored_init_params = match &mut message { ClientJsonRpcMessage::Request(req) => { let ClientRequest::InitializeRequest(init_req) = &req.request else { - return Err(unexpected_message_response("initialize request")); + return Err(unexpected_message_response("initialize request").into()); }; // Reject mismatched MCP-Protocol-Version header before binding the session to anything. validate_header_matches_init_body( @@ -1148,7 +1883,7 @@ where stored_init_params } _ => { - return Err(unexpected_message_response("initialize request")); + return Err(unexpected_message_response("initialize request").into()); } }; let service = self @@ -1219,6 +1954,11 @@ where // Stateless mode: // - on initialize: the header (if present) must match `params.protocolVersion` // - on every other request: the header must name a known version. + // + // The opt-in seam applies only here so legacy session routing and + // its error precedence remain unchanged. + validate_required_protocol_header(&self.config, &part.headers, &message)?; + let has_per_request_version = message_has_per_request_protocol_version(&message); match &message { ClientJsonRpcMessage::Request(req) => { if let ClientRequest::InitializeRequest(init_req) = &req.request { @@ -1228,18 +1968,29 @@ where Some(req.id.clone()), )?; } else { - validate_protocol_version_header(&part.headers)?; + validate_protocol_version_header(&part.headers, has_per_request_version)?; } } _ => { - validate_protocol_version_header(&part.headers)?; + validate_protocol_version_header(&part.headers, has_per_request_version)?; } } + // Validate SEP-2243 standard headers against the body + validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?; + validate_request_protocol_version_meta(&part.headers, &message)?; + validate_required_protocol_meta(&self.config, &message)?; let service = self .get_service() .map_err(internal_error_response("get service"))?; match message { ClientJsonRpcMessage::Request(mut request) => { + let negotiates_per_request = has_per_request_version + || matches!(&request.request, ClientRequest::DiscoverRequest(_)); + if negotiates_per_request { + return self + .serve_negotiated_request_directly(service, request, part) + .await; + } // Build a peer_info so context.protocol_version() works inside handlers. // serve_directly skips the handshake and receives None by default, making // protocol_version() always return None in stateless mode. We reconstruct it: @@ -1250,49 +2001,66 @@ where request.request.extensions_mut().insert(part); let (transport, mut receiver) = OneshotTransport::::new(ClientJsonRpcMessage::Request(request)); - let service = serve_directly(service, transport, peer_info); + // Give this stateless request its own cancellation token so an + // unpersisted response can cancel the in-flight handler on + // disconnect (#857). + let request_ct = CancellationToken::new(); + let service = serve_directly_with_ct( + NegotiatingStatelessHttpService(service), + transport, + peer_info, + request_ct.clone(), + ); tokio::spawn(async move { // on service created let _ = service.waiting().await; }); if self.config.json_response { - // JSON-direct mode: await the single response and return as - // application/json, eliminating SSE framing overhead. - // Allowed by MCP Streamable HTTP spec (2025-06-18). + // Prefer JSON for a terminal first message. If the handler + // emits an intermediate notification or request, preserve + // the complete message sequence by falling back to SSE. let cancel = self.config.cancellation_token.child_token(); - match tokio::select! { - res = receiver.recv() => res, - _ = cancel.cancelled() => None, - } { - Some(message) => { - tracing::trace!(?message); - let body = serde_json::to_vec(&message).map_err(|e| { - internal_error_response("serialize json response")(e) - })?; - Ok(Response::builder() - .status(http::StatusCode::OK) - .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE) - .body(Full::new(Bytes::from(body)).boxed()) - .expect("valid response")) + // Cancel the handler if the client disconnects while it is + // still producing its first message (this future is dropped + // before `receiver.recv()` completes). Disarmed once the + // handler emits anything, so a normal response is never + // cancelled. + let mut disconnect_guard = Some(request_ct.clone().drop_guard()); + let Some(message) = (tokio::select! { + res = receiver.recv() => { + if let Some(guard) = disconnect_guard.take() { + guard.disarm(); + } + res } - None => Err(internal_error_response("empty response")( + _ = cancel.cancelled() => None, + }) else { + return Err(internal_error_response("empty response")( std::io::Error::new( std::io::ErrorKind::UnexpectedEof, "no response message received from handler", ), - )), + ) + .into()); + }; + tracing::trace!(?message); + if matches!( + message, + ServerJsonRpcMessage::Response(_) | ServerJsonRpcMessage::Error(_) + ) { + let body = serde_json::to_vec(&message).map_err(|e| { + internal_error_response("serialize json response")(e) + })?; + Ok(Response::builder() + .status(http::StatusCode::OK) + .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE) + .body(Full::new(Bytes::from(body)).boxed()) + .expect("valid response")) + } else { + Ok(self.stateless_sse_response(Some(message), receiver, request_ct)) } } else { - // SSE mode (default): original behaviour preserved unchanged - let stream = ReceiverStream::new(receiver).map(|message| { - tracing::trace!(?message); - ServerSseMessage::from_message(message) - }); - Ok(sse_stream_response( - stream, - self.config.sse_keep_alive, - self.config.cancellation_token.child_token(), - )) + Ok(self.stateless_sse_response(None, receiver, request_ct)) } } ClientJsonRpcMessage::Notification(_notification) => { @@ -1305,11 +2073,14 @@ where } } - async fn handle_delete(&self, request: Request) -> Result + async fn handle_delete(&self, request: Request) -> HttpResult where B: Body + Send + 'static, B::Error: Display, { + if !is_legacy_request(None, request.headers())? { + return Ok(method_not_allowed_response()); + } // check session id let session_id = request .headers() @@ -1324,7 +2095,7 @@ where .expect("valid response")); }; // Validate MCP-Protocol-Version header (per 2025-06-18 spec) - validate_protocol_version_header(request.headers())?; + validate_protocol_version_header(request.headers(), false)?; // close session self.session_manager .close_session(&session_id) @@ -1340,7 +2111,7 @@ where Ok(accepted_response()) } - /// Build a `ClientInfo` (peer_info) for a stateless request so that + /// Build a `InitializeRequestParams` (peer_info) for a stateless request so that /// `context.protocol_version()` returns the correct value inside handlers. /// /// `serve_directly` skips the MCP handshake and accepts `peer_info = None`, @@ -1370,3 +2141,47 @@ where }) } } + +pin_project! { + /// Cancels an unpersisted stateless request when its response is dropped. + /// + /// Persisted requests keep running so another connection can resume them. + /// Without an event store, dropping the stream fires the request's + /// cancellation token. Natural completion disarms the guard. + struct CancelOnDisconnect { + #[pin] + inner: S, + ct: Option, + } + impl PinnedDrop for CancelOnDisconnect { + fn drop(this: Pin<&mut Self>) { + let this = this.project(); + if let Some(ct) = this.ct.take() { + ct.cancel(); + } + } + } +} + +impl CancelOnDisconnect { + fn new(inner: S, ct: CancellationToken) -> Self { + Self { + inner, + ct: Some(ct), + } + } +} + +impl Stream for CancelOnDisconnect { + type Item = S::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.project(); + let polled = this.inner.poll_next(cx); + if let Poll::Ready(None) = &polled { + // Ended naturally: the request completed, so don't cancel on drop. + *this.ct = None; + } + polled + } +} diff --git a/crates/rmcp/src/transport/worker.rs b/crates/rmcp/src/transport/worker.rs index 5294640..e32da7d 100644 --- a/crates/rmcp/src/transport/worker.rs +++ b/crates/rmcp/src/transport/worker.rs @@ -1,10 +1,21 @@ -use std::{borrow::Cow, time::Duration}; +use std::{ + borrow::Cow, + collections::HashMap, + sync::{ + Arc, Mutex, PoisonError, Weak, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; use tokio_util::sync::CancellationToken; use tracing::{Instrument, Level}; use super::{IntoTransport, Transport}; -use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage}; +use crate::{ + model::{CancelledNotification, JsonRpcMessage, RequestId}, + service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage}, +}; #[derive(Debug, thiserror::Error)] #[non_exhaustive] @@ -53,17 +64,114 @@ pub trait Worker: Sized + Send + 'static { fn config(&self) -> WorkerConfig { WorkerConfig::default() } + /// Return true to send this message through the separate control queue. + /// + /// Workers that opt in must read [`WorkerContext::control_from_handler_rx`] + /// and preserve any required ordering with ordinary messages. + fn is_control_message(_message: &TxJsonRpcMessage) -> bool { + false + } + /// Return true to register outgoing requests for cancellation before they enter a queue. + /// + /// Workers that opt in must honor [`WorkerSendRequest::cancellation_token`]. + fn supports_request_cancellation() -> bool { + false + } +} + +type RequestCancellations = Arc>>>; + +/// Keeps a request's cancellation token registered for a chosen lifetime. +pub(crate) struct RequestCancellationRegistration { + id: RequestId, + lifetime: CancellationToken, + cancellation: CancellationToken, + pending: RequestCancellations, +} + +impl RequestCancellationRegistration { + fn new(id: RequestId, token: CancellationToken, pending: RequestCancellations) -> Arc { + let registration = Arc::new(Self { + id: id.clone(), + cancellation: token.child_token(), + lifetime: token, + pending, + }); + registration + .pending + .lock() + .unwrap_or_else(PoisonError::into_inner) + .insert(id, Arc::downgrade(®istration)); + registration + } + + /// Return the token kept alive by this registration. + pub(crate) fn token(&self) -> CancellationToken { + self.cancellation.clone() + } + + /// Return the token cancelled when the request lifetime ends. + pub(crate) fn lifetime_token(&self) -> CancellationToken { + self.lifetime.clone() + } +} + +impl Drop for RequestCancellationRegistration { + fn drop(&mut self) { + self.lifetime.cancel(); + let mut pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner); + if pending + .get(&self.id) + .is_some_and(|current| std::ptr::eq(current.as_ptr(), self)) + { + pending.remove(&self.id); + } + } } #[non_exhaustive] pub struct WorkerSendRequest { pub message: TxJsonRpcMessage, pub responder: tokio::sync::oneshot::Sender>, + cancellation: Option>, + control_generation: u64, +} + +impl WorkerSendRequest { + /// Return the token registered before this request entered the send queue. + /// + /// This is present only for requests sent to a worker that enables + /// [`Worker::supports_request_cancellation`]. It is not sent over the wire. + /// Keep this request alive while its work is active. Cloning the token does + /// not keep its cancellation registration alive. + pub fn cancellation_token(&self) -> Option { + self.cancellation + .as_deref() + .map(RequestCancellationRegistration::token) + } + + /// Keep the same cancellation registration alive after the send completes. + #[cfg(feature = "transport-streamable-http-client")] + pub(crate) fn cancellation_registration(&self) -> Option> { + self.cancellation.clone() + } + + /// Return the local generation captured when [`Transport::send`] created its future. + /// + /// This happens before polling or queue admission. The value is not sent over + /// the wire; the worker decides whether a message from an older generation is valid. + /// It identifies the outbound send, not the session that started an inbound handler. + pub fn control_generation(&self) -> u64 { + self.control_generation + } } pub struct WorkerTransport { rx: tokio::sync::mpsc::Receiver>, send_service: tokio::sync::mpsc::Sender>, + control_send_service: tokio::sync::mpsc::Sender>, + request_cancellations: RequestCancellations, + control_generation: Arc, join_handle: Option>>>, _drop_guard: tokio_util::sync::DropGuard, ct: CancellationToken, @@ -104,11 +212,17 @@ impl WorkerTransport { let worker_name = config.name; let (to_transport_tx, from_handler_rx) = tokio::sync::mpsc::channel::>(config.channel_buffer_capacity); + let (control_to_transport_tx, control_from_handler_rx) = + tokio::sync::mpsc::channel::>(config.channel_buffer_capacity); let (to_handler_tx, from_transport_rx) = tokio::sync::mpsc::channel::>(config.channel_buffer_capacity); + let request_cancellations = RequestCancellations::default(); + let control_generation = Arc::new(AtomicU64::new(0)); let context = WorkerContext { to_handler_tx, from_handler_rx, + control_from_handler_rx, + control_generation: control_generation.clone(), cancellation_token: transport_task_ct.clone(), }; @@ -142,11 +256,32 @@ impl WorkerTransport { Self { rx: from_transport_rx, send_service: to_transport_tx, + control_send_service: control_to_transport_tx, + request_cancellations, + control_generation, join_handle: Some(join_handle), ct: transport_task_ct.clone(), _drop_guard: transport_task_ct.drop_guard(), } } + + fn cancel_request_from_notification( + &self, + notification: &::Not, + ) -> Option> { + let cancelled: CancelledNotification = notification.clone().try_into().ok()?; + let id = cancelled.params.request_id.as_ref()?; + let target = { + let pending = self + .request_cancellations + .lock() + .unwrap_or_else(PoisonError::into_inner); + pending.get(id).and_then(Weak::upgrade) + }?; + // Signal cancellation even if the control queue is full. + target.token().cancel(); + Some(target) + } } #[non_exhaustive] @@ -159,10 +294,32 @@ pub struct SendRequest { pub struct WorkerContext { pub to_handler_tx: tokio::sync::mpsc::Sender>, pub from_handler_rx: tokio::sync::mpsc::Receiver>, + /// Messages selected by [`Worker::is_control_message`]. + pub control_from_handler_rx: tokio::sync::mpsc::Receiver>, pub cancellation_token: CancellationToken, + control_generation: Arc, } impl WorkerContext { + /// Return the local generation that newly created sends will capture. + /// + /// Workers may use this value to check messages from an earlier connection + /// or session. The generic transport does not check it automatically. + pub fn control_generation(&self) -> u64 { + self.control_generation.load(Ordering::SeqCst) + } + + /// Advance the local generation, wrapping at [`u64::MAX`], and return its new value. + /// + /// Only subsequent calls to [`Transport::send`] capture the new value. + /// Advancing does not drain queues, cancel work, or reject older messages; + /// the worker is responsible for those actions. + pub fn advance_control_generation(&self) -> u64 { + self.control_generation + .fetch_add(1, Ordering::SeqCst) + .wrapping_add(1) + } + pub async fn send_to_handler( &mut self, item: RxJsonRpcMessage, @@ -190,15 +347,52 @@ impl Transport for WorkerTransport { &mut self, item: TxJsonRpcMessage, ) -> impl Future> + Send + 'static { - let tx = self.send_service.clone(); + let control_generation = self.control_generation.load(Ordering::SeqCst); + let mut cancellation_target = None; + let registration = if W::supports_request_cancellation() { + match &item { + JsonRpcMessage::Request(request) => Some(RequestCancellationRegistration::new( + request.id.clone(), + self.ct.child_token(), + self.request_cancellations.clone(), + )), + JsonRpcMessage::Notification(notification) => { + cancellation_target = + self.cancel_request_from_notification(¬ification.notification); + None + } + _ => None, + } + } else { + None + }; + let tx = if W::is_control_message(&item) { + self.control_send_service.clone() + } else { + self.send_service.clone() + }; + let cancellation_guard = registration + .as_ref() + .map(|registration| registration.lifetime_token().drop_guard()); + let target_guard = cancellation_target + .as_ref() + .map(|target| target.lifetime_token().drop_guard()); let (responder, receiver) = tokio::sync::oneshot::channel(); let request = WorkerSendRequest { message: item, responder, + cancellation: registration, + control_generation, }; async move { + // Keep the stream alive until its cancellation is handled or abandoned. + let _cancellation_target = cancellation_target; + let _target_guard = target_guard; tx.send(request).await.map_err(|_| W::err_closed())?; receiver.await.map_err(|_| W::err_closed())??; + if let Some(guard) = cancellation_guard { + let _ = guard.disarm(); + } Ok(()) } } @@ -215,3 +409,133 @@ impl Transport for WorkerTransport { } } } + +#[cfg(all(test, feature = "client"))] +mod tests { + use std::io; + + use super::*; + use crate::{model::ClientJsonRpcMessage, service::RoleClient}; + + struct TestWorker(tokio::sync::oneshot::Sender>); + + impl Worker for TestWorker { + type Error = io::Error; + type Role = RoleClient; + + fn err_closed() -> Self::Error { + io::Error::other("worker closed") + } + + fn err_join(error: tokio::task::JoinError) -> Self::Error { + io::Error::other(error) + } + + async fn run( + self, + context: WorkerContext, + ) -> Result<(), WorkerQuitReason> { + let cancellation = context.cancellation_token.clone(); + self.0 + .send(context) + .map_err(|_| WorkerQuitReason::HandlerTerminated)?; + cancellation.cancelled().await; + Ok(()) + } + + fn is_control_message(message: &ClientJsonRpcMessage) -> bool { + matches!(message, JsonRpcMessage::Notification(_)) + } + + fn supports_request_cancellation() -> bool { + true + } + } + + fn cancellation_message(id: RequestId) -> ClientJsonRpcMessage { + serde_json::from_value(serde_json::json!({ + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": { "requestId": id }, + })) + .unwrap() + } + + #[tokio::test] + async fn cancellation_matches_request_id_exactly() { + let (context_tx, context_rx) = tokio::sync::oneshot::channel(); + let mut transport = WorkerTransport::spawn(TestWorker(context_tx)); + let _context = context_rx.await.unwrap(); + let registration = RequestCancellationRegistration::new( + RequestId::Number(7), + CancellationToken::new(), + transport.request_cancellations.clone(), + ); + + drop(transport.send(cancellation_message(RequestId::String("7".into())))); + assert!(!registration.token().is_cancelled()); + + drop(transport.send(cancellation_message(RequestId::Number(7)))); + assert!(registration.token().is_cancelled()); + transport.close().await.unwrap(); + } + + #[tokio::test] + async fn abandoned_cancellation_send_ends_request_lifetime() { + let (context_tx, context_rx) = tokio::sync::oneshot::channel(); + let mut transport = WorkerTransport::spawn(TestWorker(context_tx)); + let mut context = context_rx.await.unwrap(); + + for admitted in [false, true] { + let id = RequestId::Number(7); + let registration = RequestCancellationRegistration::new( + id.clone(), + CancellationToken::new(), + transport.request_cancellations.clone(), + ); + let lifetime = registration.lifetime_token(); + let weak = Arc::downgrade(®istration); + let mut send = Box::pin(transport.send(cancellation_message(id))); + assert!(registration.token().is_cancelled()); + assert!(!lifetime.is_cancelled()); + + let queued = if admitted { + assert!(futures::poll!(send.as_mut()).is_pending()); + Some(context.control_from_handler_rx.recv().await.unwrap()) + } else { + None + }; + drop(send); + + // An open stream can still own the registration after cancellation is abandoned. + assert!(lifetime.is_cancelled()); + assert!(weak.upgrade().is_some()); + drop(registration); + assert!(weak.upgrade().is_none()); + assert!(queued.is_none_or(|request| request.responder.is_closed())); + assert!(transport.request_cancellations.lock().unwrap().is_empty()); + } + transport.close().await.unwrap(); + } + + #[test] + fn dropping_an_old_registration_preserves_a_reused_id() { + let pending = RequestCancellations::default(); + let id = RequestId::Number(7); + let old = RequestCancellationRegistration::new( + id.clone(), + CancellationToken::new(), + pending.clone(), + ); + let current = RequestCancellationRegistration::new( + id.clone(), + CancellationToken::new(), + pending.clone(), + ); + drop(old); + let registered = pending.lock().unwrap().get(&id).cloned().unwrap(); + assert!(Weak::ptr_eq(®istered, &Arc::downgrade(¤t))); + drop(current); + assert!(pending.lock().unwrap().is_empty()); + } +} diff --git a/crates/rmcp/tests/common/calculator.rs b/crates/rmcp/tests/common/calculator.rs index 22c6d38..961d977 100644 --- a/crates/rmcp/tests/common/calculator.rs +++ b/crates/rmcp/tests/common/calculator.rs @@ -2,7 +2,7 @@ use rmcp::{ ServerHandler, handler::server::{router::tool::ToolRouter, wrapper::Parameters}, - model::{ServerCapabilities, ServerInfo}, + model::{ServerCapabilities, ServerConfig}, schemars, tool, tool_router, }; #[derive(Debug, serde::Deserialize, schemars::JsonSchema)] @@ -52,8 +52,8 @@ impl Calculator { } impl ServerHandler for Calculator { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) .with_instructions("A simple calculator") } } diff --git a/crates/rmcp/tests/common/handlers.rs b/crates/rmcp/tests/common/handlers.rs index 7928fa6..7e68621 100644 --- a/crates/rmcp/tests/common/handlers.rs +++ b/crates/rmcp/tests/common/handlers.rs @@ -115,8 +115,8 @@ impl TestServer { impl ServerHandler for TestServer { #[allow(deprecated)] - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_logging().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_logging().build()) } #[allow(deprecated)] diff --git a/crates/rmcp/tests/test_auth_error_classification.rs b/crates/rmcp/tests/test_auth_error_classification.rs new file mode 100644 index 0000000..99b0508 --- /dev/null +++ b/crates/rmcp/tests/test_auth_error_classification.rs @@ -0,0 +1,85 @@ +use std::{any::TypeId, error::Error}; + +use rmcp::{ + service::ClientInitializeError, + transport::{ + AuthError, DynamicTransportError, + streamable_http_client::{AuthRequiredError, InsufficientScopeError, StreamableHttpError}, + }, +}; +use thiserror::Error; + +type TestHttpError = StreamableHttpError; + +#[derive(Debug, Error)] +#[error("outer transport wrapper")] +struct OuterError(#[source] TestHttpError); + +fn initialization_error(error: impl Error + Send + Sync + 'static) -> ClientInitializeError { + ClientInitializeError::TransportError { + error: DynamicTransportError::from_parts( + "test transport", + TypeId::of::<()>(), + Box::new(error), + ), + context: "initialize".into(), + } +} + +#[test] +fn classifies_local_authorization_required() { + let error = TestHttpError::Auth(AuthError::AuthorizationRequired); + + assert!(initialization_error(error).is_authorization_required()); +} + +#[test] +fn classifies_http_authorization_challenge() { + let error = + TestHttpError::AuthRequired(AuthRequiredError::new("Bearer realm=\"mcp\"".to_owned())); + + assert!(initialization_error(error).is_authorization_required()); +} + +#[test] +fn classifies_authorization_required_through_multiple_sources() { + let error = OuterError(TestHttpError::Auth(AuthError::AuthorizationRequired)); + + assert!(initialization_error(error).is_authorization_required()); +} + +#[test] +fn does_not_classify_unrelated_transport_errors() { + let closed = TestHttpError::TransportChannelClosed; + let refresh = TestHttpError::Auth(AuthError::TokenRefreshFailed("timeout".to_owned())); + let scope = TestHttpError::InsufficientScope(InsufficientScopeError::new( + "Bearer error=\"insufficient_scope\"".to_owned(), + Some("admin".to_owned()), + )); + + assert!(!initialization_error(closed).is_authorization_required()); + assert!(!initialization_error(refresh).is_authorization_required()); + assert!(!initialization_error(scope).is_authorization_required()); +} + +#[test] +fn does_not_classify_non_transport_initialization_errors() { + assert!(!ClientInitializeError::Cancelled.is_authorization_required()); + assert!( + !ClientInitializeError::ConnectionClosed("server closed the connection".to_owned()) + .is_authorization_required() + ); +} + +#[test] +fn http_challenge_remains_available_as_an_error_source() { + let error = + TestHttpError::AuthRequired(AuthRequiredError::new("Bearer realm=\"mcp\"".to_owned())); + + let source = error.source().expect("auth challenge should be a source"); + let challenge = source + .downcast_ref::() + .expect("source should retain the challenge type"); + + assert_eq!(challenge.www_authenticate_header, "Bearer realm=\"mcp\""); +} diff --git a/crates/rmcp/tests/test_cache_hints.rs b/crates/rmcp/tests/test_cache_hints.rs new file mode 100644 index 0000000..2b6aecd --- /dev/null +++ b/crates/rmcp/tests/test_cache_hints.rs @@ -0,0 +1,79 @@ +use rmcp::model::{CacheScope, ListToolsResult, ReadResourceResult, ResourceContents}; +use serde_json::json; + +#[test] +fn paginated_results_serialize_cache_hints_as_top_level_fields() { + let result = ListToolsResult::with_all_items(Vec::new()) + .with_ttl_ms(5_000) + .with_cache_scope(CacheScope::Private); + + let actual = serde_json::to_value(result).expect("serialize list tools result"); + + assert_eq!( + actual, + json!({ + "ttlMs": 5000, + "cacheScope": "private", + "tools": [], + "resultType": "complete" + }) + ); + assert!(actual.get("_meta").is_none()); +} + +#[test] +fn read_resource_results_serialize_cache_hints_as_top_level_fields() { + let result = + ReadResourceResult::new(vec![ResourceContents::text("hello", "file:///example.txt")]) + .with_ttl_ms(10_000) + .with_cache_scope(CacheScope::Public); + + let actual = serde_json::to_value(result).expect("serialize read resource result"); + + assert_eq!(actual["ttlMs"], 10000); + assert_eq!(actual["cacheScope"], "public"); + assert!(actual["contents"][0].get("_meta").is_none()); +} + +#[test] +fn cache_hints_are_omitted_when_absent() { + let result = ListToolsResult::with_all_items(Vec::new()); + let actual = serde_json::to_value(result).expect("serialize list tools result"); + + assert_eq!(actual, json!({ "tools": [], "resultType": "complete" })); +} + +#[test] +fn cache_hints_default_to_none_and_negative_ttl_is_normalized_to_zero() { + let absent: ListToolsResult = serde_json::from_value(json!({ + "tools": [] + })) + .expect("deserialize result without ttlMs"); + assert_eq!(absent.ttl_ms, None); + assert_eq!(absent.cache_scope, None); + + let negative: ReadResourceResult = serde_json::from_value(json!({ + "ttlMs": -42, + "cacheScope": "private", + "contents": [] + })) + .expect("deserialize result with negative ttlMs"); + assert_eq!(negative.ttl_ms, Some(0)); + assert_eq!(negative.cache_scope, Some(CacheScope::Private)); +} + +#[test] +fn cache_scope_round_trips() { + assert_eq!( + serde_json::to_value(CacheScope::Public).unwrap(), + json!("public") + ); + assert_eq!( + serde_json::to_value(CacheScope::Private).unwrap(), + json!("private") + ); + assert_eq!( + serde_json::from_value::(json!("private")).unwrap(), + CacheScope::Private + ); +} diff --git a/crates/rmcp/tests/test_cancelled_response.rs b/crates/rmcp/tests/test_cancelled_response.rs new file mode 100644 index 0000000..864932a --- /dev/null +++ b/crates/rmcp/tests/test_cancelled_response.rs @@ -0,0 +1,633 @@ +//! A receiver SHOULD NOT send a response for a request it has already been told +//! to cancel. This drives a real stdio server with raw JSON-RPC: the tool blocks +//! until the request is cancelled, so its result is only produced *after* the +//! cancellation — the service loop must drop it rather than write it to the wire. + +#[cfg(all(feature = "client", not(feature = "local")))] +use std::sync::Arc; +use std::{collections::BTreeSet, process::Stdio, time::Duration}; + +#[cfg(all(feature = "client", not(feature = "local")))] +use rmcp::{ + ClientHandler, RoleClient, + model::{ + CancelledNotification, CancelledNotificationParam, ClientJsonRpcMessage, ClientRequest, + ClientResult, ElicitRequest, ElicitRequestParams, ElicitResult, ElicitationAction, + ElicitationSchema, PingRequest, RequestId, ServerJsonRpcMessage, ServerNotification, + ServerRequest, ServerResult, + }, + service::{PeerRequestOptions, QuitReason, serve_directly}, + transport::{IntoTransport, Transport}, +}; +use rmcp::{ + ErrorData as McpError, RoleServer, ServerHandler, ServiceExt, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ServerCapabilities, + ServerConfig, + }, + service::RequestContext, +}; +use serde_json::{Value, json}; +use tokio::{ + io::{AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader}, + process::{Child, Command}, +}; + +const HELPER_ENV: &str = "RMCP_CANCELLED_RESPONSE_HELPER"; +const READ_TIMEOUT: Duration = Duration::from_secs(10); +const STRING_REQUEST_ID: &str = "tool-request-2"; + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cancelled_request_receives_no_response() -> anyhow::Result<()> { + let mut child = spawn_helper(); + let mut writer = child.stdin.take().expect("helper stdin"); + let stdout = child.stdout.take().expect("helper stdout"); + let mut reader = BufReader::new(stdout); + + send_json( + &mut writer, + &json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": { "name": "raw-test-client", "version": "0.0.0" } + } + }), + ) + .await?; + collect_ids_until(&mut reader, "1", READ_TIMEOUT).await?; + send_json( + &mut writer, + &json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }), + ) + .await?; + + // Start a request that blocks until cancelled, then cancel it. Its response is + // produced only after the cancellation arrives, so it must be suppressed. + send_json( + &mut writer, + &json!({ + "jsonrpc": "2.0", + "id": STRING_REQUEST_ID, + "method": "tools/call", + "params": { "name": "wait-for-cancel", "arguments": {} } + }), + ) + .await?; + send_json( + &mut writer, + &json!({ + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": { "requestId": "unrelated-request" } + }), + ) + .await?; + send_json( + &mut writer, + &json!({ + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": { "requestId": STRING_REQUEST_ID } + }), + ) + .await?; + // A ping proves the server is alive past the cancellation, so the absence of + // a cancelled response is genuine suppression rather than a dead connection. + send_json( + &mut writer, + &json!({ "jsonrpc": "2.0", "id": 3, "method": "ping" }), + ) + .await?; + + let seen = collect_ids_until(&mut reader, "3", READ_TIMEOUT).await?; + assert!(seen.contains("3")); + assert!(!seen.contains(STRING_REQUEST_ID)); + + send_json( + &mut writer, + &json!({ + "jsonrpc": "2.0", + "id": 4, + "method": "tools/call", + "params": { "name": "complete-unless-cancelled", "arguments": {} } + }), + ) + .await?; + send_json( + &mut writer, + &json!({ + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": { "requestId": "4" } + }), + ) + .await?; + send_json( + &mut writer, + &json!({ "jsonrpc": "2.0", "id": 5, "method": "ping" }), + ) + .await?; + + let seen = collect_ids_until(&mut reader, "5", READ_TIMEOUT).await?; + assert!(seen.contains("4"), "string and numeric IDs must not alias"); + assert!(seen.contains("5")); + + drop(writer); + wait_for_child(&mut child).await; + Ok(()) +} + +struct WaitForCancelServer; + +impl ServerHandler for WaitForCancelServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + if request.name == "complete-unless-cancelled" { + tokio::select! { + _ = context.ct.cancelled() => {} + _ = tokio::time::sleep(Duration::from_millis(100)) => {} + } + return Ok(CallToolResult::success(vec![ContentBlock::text("completed")]).into()); + } + context.ct.cancelled().await; + Ok(CallToolResult::success(vec![ContentBlock::text("late response")]).into()) + } +} + +#[cfg(all(feature = "client", not(feature = "local")))] +#[derive(Clone)] +struct WaitForReverseCancelClient { + events: tokio::sync::mpsc::UnboundedSender<&'static str>, + finish: Option>, + return_error: bool, +} + +#[cfg(all(feature = "client", not(feature = "local")))] +impl ClientHandler for WaitForReverseCancelClient { + async fn create_elicitation( + &self, + _request: ElicitRequestParams, + context: rmcp::service::RequestContext, + ) -> Result { + self.events.send("started").expect("test is listening"); + context.ct.cancelled().await; + self.events.send("cancelled").expect("test is listening"); + if let Some(finish) = &self.finish { + finish.notified().await; + } + if self.return_error { + return Err(McpError::internal_error("late handler error", None)); + } + Ok(ElicitResult::new(ElicitationAction::Decline)) + } +} + +#[cfg(all(feature = "client", not(feature = "local")))] +#[rstest::rstest] +#[case::initialize("initialize")] +#[case::direct_legacy("direct-legacy")] +#[case::direct_unknown("direct-unknown")] +#[tokio::test] +async fn peer_cancels_reverse_request_without_cancelling_outbound_request( + #[case] startup: &str, + #[values(false, true)] equal_ids: bool, +) -> anyhow::Result<()> { + tokio::time::timeout(READ_TIMEOUT, reverse_cancellation(startup, equal_ids)).await??; + Ok(()) +} + +#[cfg(all(feature = "client", not(feature = "local")))] +async fn reverse_cancellation(startup: &str, equal_ids: bool) -> anyhow::Result<()> { + use rmcp::model::ProtocolVersion; + let (client_transport, server_transport) = tokio::io::duplex(4096); + let (events_tx, mut events_rx) = tokio::sync::mpsc::unbounded_channel(); + let handler = WaitForReverseCancelClient { + events: events_tx, + finish: None, + return_error: false, + }; + let mut server = IntoTransport::::into_transport(server_transport); + let mut info = ServerConfig::default(); + info.protocol_version = ProtocolVersion::V_2025_11_25; + let client = if startup == "initialize" { + let (client, handshake) = tokio::join!(handler.serve(client_transport), async { + let Some(ClientJsonRpcMessage::Request(request)) = server.receive().await else { + panic!("expected initialize request"); + }; + assert!(matches!( + request.request, + ClientRequest::InitializeRequest(_) + )); + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult(info.clone()), + request.id, + )) + .await?; + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + anyhow::Ok(()) + }); + handshake?; + client? + } else { + let peer_info = (startup == "direct-legacy").then(|| info.into()); + serve_directly::(handler, client_transport, peer_info) + }; + let mut outbound = client + .send_cancellable_request( + ClientRequest::PingRequest(PingRequest { + method: Default::default(), + extensions: Default::default(), + }), + PeerRequestOptions::no_options(), + ) + .await?; + let Some(ClientJsonRpcMessage::Request(outbound_request)) = server.receive().await else { + panic!("expected outbound ping request"); + }; + let id = if equal_ids { + outbound_request.id.clone() + } else { + RequestId::String("elicitation-1".into()) + }; + let mut request = reverse_elicitation("elicitation-1"); + if let ServerJsonRpcMessage::Request(request) = &mut request { + request.id = id.clone(); + } + server.send(request).await?; + assert_eq!(events_rx.recv().await, Some("started")); + let unknown = if equal_ids { + RequestId::String(id.to_string().into()) + } else { + outbound_request.id.clone() + }; + for cancel_id in [unknown, RequestId::String("unknown".into())] { + server.send(cancel_notification(cancel_id)).await?; + } + exchange_ping(&mut server, "before-cancel").await?; + assert!(matches!( + events_rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + assert!(matches!( + outbound.rx.try_recv(), + Err(tokio::sync::oneshot::error::TryRecvError::Empty) + )); + + // A cancellation can arrive after its response was already delivered. + server + .send(cancel_notification(RequestId::String( + "before-cancel".into(), + ))) + .await?; + server.send(cancel_notification(id)).await?; + assert_eq!(events_rx.recv().await, Some("cancelled")); + exchange_ping(&mut server, "after-cancel").await?; + assert!(matches!( + outbound.rx.try_recv(), + Err(tokio::sync::oneshot::error::TryRecvError::Empty) + )); + server + .send(ServerJsonRpcMessage::response( + ServerResult::empty(()), + outbound_request.id, + )) + .await?; + assert!(matches!( + outbound.await_response().await?, + ServerResult::EmptyResult(_) + )); + + client.cancel().await?; + assert!( + server.receive().await.is_none(), + "no late response after handler drain" + ); + Ok(()) +} + +#[cfg(all(feature = "client", not(feature = "local")))] +fn cancel_notification(id: RequestId) -> ServerJsonRpcMessage { + ServerJsonRpcMessage::notification( + CancelledNotification::new(CancelledNotificationParam::new( + Some(id), + Some("request cancelled".to_owned()), + )) + .into(), + ) +} + +#[cfg(all(feature = "client", not(feature = "local")))] +async fn exchange_ping(server: &mut impl Transport, id: &str) -> anyhow::Result<()> { + server + .send(ServerJsonRpcMessage::request( + PingRequest { + method: Default::default(), + extensions: Default::default(), + } + .into(), + RequestId::String(id.into()), + )) + .await?; + let Some(ClientJsonRpcMessage::Response(response)) = server.receive().await else { + panic!("expected ping response"); + }; + assert_eq!(response.id, RequestId::String(id.into())); + assert!(matches!(response.result, ClientResult::EmptyResult(_))); + Ok(()) +} + +#[cfg(all(feature = "client", not(feature = "local")))] +struct ChannelClientTransport { + incoming: tokio::sync::mpsc::UnboundedReceiver, + outgoing: tokio::sync::mpsc::UnboundedSender, + eof_observed: Arc, +} + +#[cfg(all(feature = "client", not(feature = "local")))] +impl Transport for ChannelClientTransport { + type Error = std::io::Error; + + fn send( + &mut self, + item: ClientJsonRpcMessage, + ) -> impl Future> + Send + 'static { + let outgoing = self.outgoing.clone(); + async move { + outgoing.send(item).map_err(|_| { + std::io::Error::new(std::io::ErrorKind::BrokenPipe, "test receiver closed") + }) + } + } + + async fn receive(&mut self) -> Option { + let message = self.incoming.recv().await; + if message.is_none() { + self.eof_observed.notify_one(); + } + message + } + + async fn close(&mut self) -> Result<(), Self::Error> { + Ok(()) + } +} + +#[cfg(all(feature = "client", not(feature = "local")))] +fn reverse_elicitation(id: &str) -> ServerJsonRpcMessage { + ServerJsonRpcMessage::request( + ServerRequest::ElicitRequest(ElicitRequest::new( + ElicitRequestParams::FormElicitationParams { + meta: None, + message: "Continue the operation?".to_owned(), + requested_schema: ElicitationSchema::builder() + .build() + .expect("empty elicitation schema is valid"), + }, + )), + RequestId::String(id.into()), + ) +} + +#[cfg(all(feature = "client", not(feature = "local")))] +#[tokio::test] +async fn graceful_close_drains_uncancelled_reverse_response() -> anyhow::Result<()> { + const REVERSE_REQUEST_ID: &str = "elicitation-close-drain"; + + let (to_client, incoming) = tokio::sync::mpsc::unbounded_channel(); + let (outgoing, mut from_client) = tokio::sync::mpsc::unbounded_channel(); + let eof_observed = Arc::new(tokio::sync::Notify::new()); + let finish = Arc::new(tokio::sync::Notify::new()); + let (events_tx, mut events_rx) = tokio::sync::mpsc::unbounded_channel(); + let mut client = serve_directly::( + WaitForReverseCancelClient { + events: events_tx, + finish: Some(finish.clone()), + return_error: false, + }, + ChannelClientTransport { + incoming, + outgoing, + eof_observed, + }, + None, + ); + + to_client.send(reverse_elicitation(REVERSE_REQUEST_ID))?; + assert_eq!(events_rx.recv().await, Some("started")); + + let (close_result, cancellation_event) = tokio::join!(client.close(), async { + let event = events_rx.recv().await; + finish.notify_one(); + event + }); + assert!(matches!(close_result?, QuitReason::Cancelled)); + assert_eq!(cancellation_event, Some("cancelled")); + + let Some(ClientJsonRpcMessage::Response(response)) = + tokio::time::timeout(READ_TIMEOUT, from_client.recv()).await? + else { + panic!("expected in-flight response during graceful close"); + }; + assert_eq!(response.id, RequestId::String(REVERSE_REQUEST_ID.into())); + Ok(()) +} + +#[cfg(all(feature = "client", not(feature = "local")))] +#[tokio::test] +async fn cancelled_reverse_request_stays_suppressed_during_eof_drain() -> anyhow::Result<()> { + for return_error in [false, true] { + tokio::time::timeout( + READ_TIMEOUT, + cancelled_reverse_response_at_eof(return_error), + ) + .await??; + } + Ok(()) +} + +#[cfg(all(feature = "client", not(feature = "local")))] +async fn cancelled_reverse_response_at_eof(return_error: bool) -> anyhow::Result<()> { + const REVERSE_REQUEST_ID: &str = "elicitation-before-eof"; + + let (to_client, incoming) = tokio::sync::mpsc::unbounded_channel(); + let (outgoing, mut from_client) = tokio::sync::mpsc::unbounded_channel(); + let eof_observed = Arc::new(tokio::sync::Notify::new()); + let finish = Arc::new(tokio::sync::Notify::new()); + let (events_tx, mut events_rx) = tokio::sync::mpsc::unbounded_channel(); + let client = serve_directly::( + WaitForReverseCancelClient { + events: events_tx, + finish: Some(finish.clone()), + return_error, + }, + ChannelClientTransport { + incoming, + outgoing, + eof_observed: eof_observed.clone(), + }, + None, + ); + + to_client.send(reverse_elicitation(REVERSE_REQUEST_ID))?; + assert_eq!( + tokio::time::timeout(READ_TIMEOUT, events_rx.recv()).await?, + Some("started") + ); + + to_client.send(ServerJsonRpcMessage::notification( + ServerNotification::CancelledNotification(CancelledNotification::new( + CancelledNotificationParam::new( + Some(RequestId::String(REVERSE_REQUEST_ID.into())), + Some("user cancelled".to_owned()), + ), + )), + ))?; + assert_eq!( + tokio::time::timeout(READ_TIMEOUT, events_rx.recv()).await?, + Some("cancelled") + ); + + let eof = eof_observed.notified(); + drop(to_client); + tokio::time::timeout(READ_TIMEOUT, eof).await?; + finish.notify_one(); + + assert!( + matches!( + tokio::time::timeout(READ_TIMEOUT, client.waiting()).await??, + QuitReason::Closed + ), + "client must close after input EOF" + ); + assert!( + from_client.recv().await.is_none(), + "cancelled reverse request must not be sent during EOF drain (return_error={return_error})" + ); + Ok(()) +} + +#[tokio::test] +async fn cancelled_response_helper() -> anyhow::Result<()> { + if std::env::var(HELPER_ENV).as_deref() != Ok("1") { + return Ok(()); + } + run_helper_server().await?; + Ok(()) +} + +#[cfg(feature = "local")] +async fn run_helper_server() -> anyhow::Result<()> { + tokio::task::LocalSet::new() + .run_until(serve_helper_stdio()) + .await +} + +#[cfg(not(feature = "local"))] +async fn run_helper_server() -> anyhow::Result<()> { + serve_helper_stdio().await +} + +async fn serve_helper_stdio() -> anyhow::Result<()> { + let server = WaitForCancelServer.serve(rmcp::transport::stdio()).await?; + server.waiting().await?; + Ok(()) +} + +fn spawn_helper() -> Child { + let exe = std::env::current_exe().expect("current test exe"); + Command::new(exe) + .arg("--exact") + .arg("cancelled_response_helper") + .arg("--quiet") + .arg("--nocapture") + .arg("--test-threads") + .arg("1") + .env(HELPER_ENV, "1") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .kill_on_drop(true) + .spawn() + .expect("spawn helper") +} + +async fn wait_for_child(child: &mut Child) { + let _ = tokio::time::timeout(Duration::from_secs(2), child.wait()).await; + if child.id().is_some() { + let _ = child.kill().await; + } +} + +async fn send_json(writer: &mut W, message: &Value) -> anyhow::Result<()> +where + W: AsyncWrite + Unpin, +{ + let serialized = serde_json::to_string(message)?; + writer.write_all(serialized.as_bytes()).await?; + writer.write_all(b"\n").await?; + writer.flush().await?; + Ok(()) +} + +/// Read response lines, collecting every message id seen, until `stop_id` is seen +/// (then a short grace read to catch any straggler) or the timeout elapses. +async fn collect_ids_until( + reader: &mut BufReader, + stop_id: &str, + timeout: Duration, +) -> anyhow::Result> +where + R: tokio::io::AsyncRead + Unpin, +{ + let mut seen = BTreeSet::new(); + let mut deadline = tokio::time::Instant::now() + timeout; + loop { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + break; + } + let mut line = String::new(); + let Ok(read_result) = tokio::time::timeout(remaining, reader.read_line(&mut line)).await + else { + break; + }; + if read_result? == 0 { + break; + } + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let Ok(value) = serde_json::from_str::(trimmed) else { + continue; + }; + let id = match value.get("id") { + Some(Value::String(id)) => Some(id.clone()), + Some(Value::Number(id)) => Some(id.to_string()), + _ => None, + }; + if let Some(id) = id { + seen.insert(id.clone()); + if id == stop_id { + // Give any late (incorrectly-sent) response a brief window to arrive. + deadline = tokio::time::Instant::now() + Duration::from_millis(300); + } + } + } + Ok(seen) +} diff --git a/crates/rmcp/tests/test_client_credentials.rs b/crates/rmcp/tests/test_client_credentials.rs index b2698d1..f21c858 100644 --- a/crates/rmcp/tests/test_client_credentials.rs +++ b/crates/rmcp/tests/test_client_credentials.rs @@ -37,14 +37,29 @@ async fn resource_metadata_handler(req: Request) -> Result, async fn auth_server_metadata_handler(req: Request) -> Result, Infallible> { let host = req.headers().get("host").unwrap().to_str().unwrap(); let base_url = format!("http://{}", host); - Ok(json_response(serde_json::json!({ - "issuer": base_url, - "authorization_endpoint": format!("{}/authorize", base_url), - "token_endpoint": format!("{}/token", base_url), + Ok(auth_server_metadata_response(&base_url, &base_url)) +} + +async fn path_inserted_auth_server_metadata_handler( + req: Request, +) -> Result, Infallible> { + let host = req.headers().get("host").unwrap().to_str().unwrap(); + let base_url = format!("http://{}", host); + Ok(auth_server_metadata_response( + &format!("{}/mcp", base_url), + &base_url, + )) +} + +fn auth_server_metadata_response(issuer: &str, endpoint_base_url: &str) -> Response { + json_response(serde_json::json!({ + "issuer": issuer, + "authorization_endpoint": format!("{}/authorize", endpoint_base_url), + "token_endpoint": format!("{}/token", endpoint_base_url), "token_endpoint_auth_methods_supported": ["client_secret_post", "client_secret_basic"], "grant_types_supported": ["client_credentials"], "scopes_supported": ["read", "write"] - }))) + })) } async fn token_handler(req: Request) -> Result, Infallible> { @@ -144,7 +159,7 @@ async fn start_path_insert_metadata_server() -> (String, SocketAddr) { let app = Router::new() .route( "/.well-known/oauth-authorization-server/mcp", - get(auth_server_metadata_handler), + get(path_inserted_auth_server_metadata_handler), ) .route("/token", post(token_handler)); diff --git a/crates/rmcp/tests/test_client_initialization.rs b/crates/rmcp/tests/test_client_initialization.rs index f51b33e..0ca0182 100644 --- a/crates/rmcp/tests/test_client_initialization.rs +++ b/crates/rmcp/tests/test_client_initialization.rs @@ -9,11 +9,110 @@ use common::handlers::TestClientHandler; use rmcp::{ ServiceExt, model::{ - ErrorCode, ErrorData, JsonRpcError, JsonRpcVersion2_0, RequestId, ServerJsonRpcMessage, + ClientJsonRpcMessage, ErrorCode, ErrorData, InitializeResult, JsonRpcError, + JsonRpcVersion2_0, RequestId, ServerCapabilities, ServerJsonRpcMessage, ServerResult, }, transport::{IntoTransport, Transport}, }; +fn stringify_numeric_id(id: RequestId) -> RequestId { + let RequestId::Number(id) = id else { + panic!("expected a numeric request ID"); + }; + RequestId::String(id.to_string().into()) +} + +#[tokio::test] +async fn client_initialization_accepts_stringified_numeric_response_id() { + let (server_transport, client_transport) = tokio::io::duplex(1024); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(request) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected initialize request"); + }; + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult( + InitializeResult::new(ServerCapabilities::default()), + ), + stringify_numeric_id(request.id), + )) + .await + .expect("send initialize response"); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + }); + + let client = TestClientHandler::new(true, true) + .serve(client_transport) + .await + .expect("client should accept stringified initialize response ID"); + assert!( + client + .peer_info() + .expect("peer info should be retained") + .server_info + .is_some(), + "initialize always provides a server implementation identity" + ); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn client_correlates_stringified_numeric_response_id() { + let (server_transport, client_transport) = tokio::io::duplex(1024); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(initialize) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected initialize request"); + }; + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult( + InitializeResult::new(ServerCapabilities::default()), + ), + initialize.id, + )) + .await + .expect("send initialize response"); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + + let ClientJsonRpcMessage::Request(request) = + server.receive().await.expect("expected tools/list request") + else { + panic!("expected tools/list request"); + }; + server + .send(ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(Default::default()), + stringify_numeric_id(request.id), + )) + .await + .expect("send tools/list response"); + }); + + let client = TestClientHandler::new(true, true) + .serve(client_transport) + .await + .expect("initialize client"); + client + .list_tools(None) + .await + .expect("client should correlate stringified response ID"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + #[tokio::test] async fn test_client_init_handles_jsonrpc_error() { let (server_transport, client_transport) = tokio::io::duplex(1024); @@ -26,11 +125,18 @@ async fn test_client_init_handles_jsonrpc_error() { }); tokio::spawn(async move { - let _init_request = server.receive().await; + let request = server.receive().await; + // Echo the request's own id back on the error so it correlates: an + // uncorrelated id would surface as `UncorrelatedErrorResponse` + // instead of the `JsonRpcError` this test exercises. + let request_id = request + .and_then(|message| message.into_request()) + .map(|(_, id)| id) + .expect("client sent an initialize request"); let error_msg = ServerJsonRpcMessage::Error(JsonRpcError { jsonrpc: JsonRpcVersion2_0, - id: Some(RequestId::Number(1)), + id: Some(request_id), error: ErrorData { code: ErrorCode(-32600), message: Cow::Borrowed("Invalid Request"), diff --git a/crates/rmcp/tests/test_client_lifecycle_modes.rs b/crates/rmcp/tests/test_client_lifecycle_modes.rs new file mode 100644 index 0000000..5fa2cd1 --- /dev/null +++ b/crates/rmcp/tests/test_client_lifecycle_modes.rs @@ -0,0 +1,791 @@ +#![cfg(all(feature = "client", feature = "server", not(feature = "local")))] + +use rmcp::{ + ClientHandler, ClientLifecycleMode, ClientServiceExt, ServerHandler, ServiceExt, + model::{ + ClientJsonRpcMessage, ClientRequest, DiscoverResult, ErrorCode, ErrorData, GetMeta, + Implementation, InitializeResult, ProtocolVersion, RequestId, ServerCapabilities, + ServerJsonRpcMessage, ServerResult, + }, + service::PeerRequestOptions, + transport::{IntoTransport, Transport}, +}; + +#[derive(Clone, Default)] +struct DiscoverClient; + +impl ClientHandler for DiscoverClient {} + +#[derive(Clone, Default)] +struct StatelessServer; + +impl ServerHandler for StatelessServer {} + +#[tokio::test] +async fn discover_startup_accepts_stringified_numeric_response_id() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(request) = + server.receive().await.expect("expected discover request") + else { + panic!("expected discover request"); + }; + let RequestId::Number(response_id) = request.id else { + panic!("expected a numeric request ID"); + }; + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult( + DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + ServerCapabilities::default(), + ) + .with_server_info(Implementation::new("discover-server", "1.0.0")), + ), + RequestId::String(response_id.to_string().into()), + )) + .await + .expect("send discover response"); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("client should accept stringified discover response ID"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +#[allow(deprecated)] +async fn discover_startup_accepts_missing_optional_server_info() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let (rejection_observed_tx, rejection_observed_rx) = tokio::sync::oneshot::channel(); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(discover_request) = + server.receive().await.expect("expected discover request") + else { + panic!("expected discover request"); + }; + let mut result = DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + ServerCapabilities::builder().enable_tools().build(), + ); + result.instructions = Some("discovery instructions".into()); + result.meta = Some(rmcp::model::MetaObject::new()); + result + .meta + .as_mut() + .expect("metadata") + .0 + .insert("example.test/key".into(), serde_json::json!(7)); + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult(result), + discover_request.id, + )) + .await + .expect("send discover response"); + + server + .send(ServerJsonRpcMessage::request( + rmcp::model::ServerRequest::CreateMessageRequest( + rmcp::model::CreateMessageRequest::new( + rmcp::model::CreateMessageRequestParams::new( + vec![rmcp::model::SamplingMessage::user_text("unsolicited")], + 16, + ), + ), + ), + RequestId::Number(99), + )) + .await + .expect("send unsolicited server request"); + let Some(ClientJsonRpcMessage::Error(error)) = server.receive().await else { + panic!("expected unsolicited server request to be rejected"); + }; + assert_eq!(error.error.code, ErrorCode::INVALID_PARAMS); + rejection_observed_tx + .send(()) + .expect("signal observed rejection"); + + let ClientJsonRpcMessage::Request(request) = + server.receive().await.expect("expected normal request") + else { + panic!("expected normal request"); + }; + assert_eq!( + request.request.get_meta().protocol_version(), + Some(ProtocolVersion::V_2026_07_28) + ); + server + .send(ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(Default::default()), + request.id, + )) + .await + .expect("send list tools response"); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("missing optional server info should not fail discovery"); + let peer_info = client.peer_info().expect("peer info should be retained"); + assert_eq!(peer_info.protocol_version, ProtocolVersion::V_2026_07_28); + assert!(peer_info.capabilities.tools.is_some()); + assert_eq!(peer_info.server_info, None); + assert_eq!( + peer_info.instructions.as_deref(), + Some("discovery instructions") + ); + assert_eq!( + peer_info + .meta + .as_ref() + .and_then(|meta| meta.0.get("example.test/key")), + Some(&serde_json::json!(7)) + ); + rejection_observed_rx + .await + .expect("server should observe association rejection"); + client.list_tools(None).await.expect("list tools"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn high_level_server_accepts_discover_startup_without_initialize() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server_task = tokio::spawn(async move { + StatelessServer + .serve(server_transport) + .await + .expect("server should accept discover") + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("discover client should start"); + client.list_tools(None).await.expect("list tools"); + client.cancel().await.expect("cancel client"); + let server = server_task.await.expect("server task"); + server.cancel().await.expect("cancel server"); +} + +#[tokio::test] +async fn discover_startup_omits_initialize() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(request) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + assert!(matches!(request.request, ClientRequest::DiscoverRequest(_))); + let meta = request.request.get_meta(); + assert_eq!(meta.protocol_version(), Some(ProtocolVersion::V_2026_07_28)); + assert!(meta.client_info().is_some()); + assert!(meta.client_capabilities().is_some()); + + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult( + DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + ServerCapabilities::default(), + ) + .with_server_info(Implementation::new("discover-server", "1.0.0")), + ), + request.id, + )) + .await + .expect("send discover response"); + + let ClientJsonRpcMessage::Request(request) = + server.receive().await.expect("expected normal request") + else { + panic!("expected request"); + }; + assert!(!matches!( + request.request, + ClientRequest::InitializeRequest(_) + )); + let meta = request.request.get_meta(); + assert_eq!(meta.protocol_version(), Some(ProtocolVersion::V_2025_11_25)); + assert!(meta.client_info().is_some()); + assert!(meta.client_capabilities().is_some()); + assert_eq!( + meta.get("example.test/extension"), + Some(&serde_json::json!(7)) + ); + server + .send(ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(Default::default()), + request.id, + )) + .await + .expect("send tools response"); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("discover client should start"); + let mut caller_meta = rmcp::model::RequestMetaObject::new(); + caller_meta.insert("example.test/extension".into(), serde_json::json!(7)); + caller_meta.set_protocol_version(ProtocolVersion::V_2025_11_25); + client + .send_request_with_option( + ClientRequest::ListToolsRequest(rmcp::model::ListToolsRequest { + method: Default::default(), + params: None, + extensions: Default::default(), + }), + PeerRequestOptions::default().with_meta(caller_meta), + ) + .await + .expect("send list tools") + .await_response() + .await + .expect("list tools response"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn auto_startup_falls_back_after_discover_method_not_found() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + assert!(matches!( + discover.request, + ClientRequest::DiscoverRequest(_) + )); + server + .send(ServerJsonRpcMessage::error( + ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "Method not found", None), + Some(discover.id), + )) + .await + .expect("send method-not-found"); + + let ClientJsonRpcMessage::Request(initialize) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected request"); + }; + assert!(matches!( + initialize.request, + ClientRequest::InitializeRequest(_) + )); + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult( + InitializeResult::new(ServerCapabilities::default()), + ), + initialize.id, + )) + .await + .expect("send initialize response"); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await + .expect("auto client should fall back"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +/// Drives an `Auto` client through a single `server/discover` probe and asserts +/// the legacy fallback decision against the response the server sends back. +/// +/// When `expect_fallback` is set, the server also accepts the subsequent +/// `initialize` request and the client is expected to connect. Otherwise the +/// client must surface the discover error without sending `initialize`, and the +/// server's next receive must not be an initialize request. +async fn run_auto_discover_response_scenario(error: ErrorData, expect_fallback: bool) { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + assert!(matches!( + discover.request, + ClientRequest::DiscoverRequest(_) + )); + server + .send(ServerJsonRpcMessage::error(error, Some(discover.id))) + .await + .expect("send discover error response"); + + if expect_fallback { + let ClientJsonRpcMessage::Request(initialize) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected request"); + }; + assert!(matches!( + initialize.request, + ClientRequest::InitializeRequest(_) + )); + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult(InitializeResult::new( + ServerCapabilities::default(), + )), + initialize.id, + )) + .await + .expect("send initialize response"); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + } else { + // The client must surface the error without falling back, so no + // initialize request should follow. The transport closes when the + // failed client is dropped. + if let Some(ClientJsonRpcMessage::Request(request)) = server.receive().await { + panic!( + "client fell back to {:?} but should have surfaced the modern error", + request.request + ); + } + } + }); + + let client_result = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await; + + if expect_fallback { + let client = client_result.expect("auto client should fall back to initialize"); + client.cancel().await.expect("cancel client"); + } else { + assert!( + client_result.is_err(), + "modern error should surface without legacy fallback" + ); + } + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn auto_startup_falls_back_after_discover_invalid_request() { + // Legacy servers commonly reject an unknown pre-initialize request with + // `-32600` (e.g. a session middleware that requires `initialize` first). + run_auto_discover_response_scenario( + ErrorData::new(ErrorCode::INVALID_REQUEST, "Bad Request", None), + true, + ) + .await; +} + +#[tokio::test] +async fn auto_startup_falls_back_after_discover_invalid_params() { + // `-32602` is explicitly called out by the specification as an + // implementation-defined response legacy servers use for unknown requests. + run_auto_discover_response_scenario( + ErrorData::new(ErrorCode::INVALID_PARAMS, "Invalid params", None), + true, + ) + .await; +} + +#[tokio::test] +async fn auto_startup_does_not_fall_back_for_missing_required_capability() { + // A `MISSING_REQUIRED_CLIENT_CAPABILITY` response identifies a modern + // server; falling back to `initialize` would not address it. + run_auto_discover_response_scenario( + ErrorData::new( + ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY, + "Missing required client capability", + None, + ), + false, + ) + .await; +} + +#[tokio::test] +async fn auto_startup_does_not_fall_back_for_header_mismatch() { + // A `HEADER_MISMATCH` response identifies a modern server performing + // header validation; falling back to `initialize` would not address it. + run_auto_discover_response_scenario( + ErrorData::new(ErrorCode::HEADER_MISMATCH, "Header mismatch", None), + false, + ) + .await; +} + +#[tokio::test] +async fn auto_startup_preserves_both_errors_when_fallback_also_fails() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + // Discover: legacy rejection. + let ClientJsonRpcMessage::Request(discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + server + .send(ServerJsonRpcMessage::error( + ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "not found", None), + Some(discover.id), + )) + .await + .expect("send discover error"); + + // Initialize: close the transport instead of responding. + let ClientJsonRpcMessage::Request(_) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected request"); + }; + drop(server); // closes the transport + }); + + let result = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await; + + let err = result.err().expect("both phases should fail"); + match err { + rmcp::service::ClientInitializeError::LegacyFallbackFailed { discover, fallback } => { + assert!( + matches!( + *discover, + rmcp::service::ClientInitializeError::JsonRpcError(_) + ), + "discover phase should be a JsonRpcError, got {discover:?}" + ); + assert!( + matches!( + *fallback, + rmcp::service::ClientInitializeError::ConnectionClosed(_) + ), + "fallback phase should be ConnectionClosed, got {fallback:?}" + ); + } + other => panic!("expected LegacyFallbackFailed, got {other:?}"), + } + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn auto_startup_surfaces_uncorrelated_error_id() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(_) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + // Respond with an error carrying a different id. + server + .send(ServerJsonRpcMessage::error( + ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "not found", None), + Some(RequestId::Number(999)), + )) + .await + .expect("send mismatched-id error"); + }); + + let result = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await; + + assert!( + matches!( + result, + Err(rmcp::service::ClientInitializeError::UncorrelatedErrorResponse { .. }) + ), + "mismatched id should surface as UncorrelatedErrorResponse, got {:?}", + result.as_ref().err() + ); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn auto_startup_falls_back_for_absent_error_id() { + // Spec: error id is optional; a server that cannot read the request id + // omits it. The error is still a response to our request, so it should + // trigger legacy fallback. + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(_discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + server + .send(ServerJsonRpcMessage::error( + ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "not found", None), + None, + )) + .await + .expect("send no-id error"); + + // Client should fall back to initialize. + let ClientJsonRpcMessage::Request(initialize) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected request"); + }; + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult( + InitializeResult::new(ServerCapabilities::default()), + ), + initialize.id, + )) + .await + .expect("send initialize response"); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await + .expect("absent-id error should trigger legacy fallback"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn discover_mode_surfaces_legacy_error() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + server + .send(ServerJsonRpcMessage::error( + ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "not found", None), + Some(discover.id), + )) + .await + .expect("send discover error"); + }); + + let result = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await; + + assert!( + matches!( + result, + Err(rmcp::service::ClientInitializeError::JsonRpcError(_)) + ), + "Discover mode should surface a legacy error, not fall back, got {:?}", + result.as_ref().err() + ); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn discover_startup_retries_a_mutually_supported_version() { + let unsupported: ProtocolVersion = + serde_json::from_value(serde_json::json!("2099-01-01")).unwrap(); + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(first) = + server.receive().await.expect("expected first discover") + else { + panic!("expected request"); + }; + assert_eq!( + first.request.get_meta().protocol_version(), + Some(unsupported.clone()) + ); + server + .send(ServerJsonRpcMessage::error( + ErrorData::unsupported_protocol_version( + unsupported, + &[ProtocolVersion::V_2026_07_28], + ), + Some(first.id), + )) + .await + .expect("send unsupported error"); + + let ClientJsonRpcMessage::Request(second) = + server.receive().await.expect("expected retry discover") + else { + panic!("expected request"); + }; + assert_eq!( + second.request.get_meta().protocol_version(), + Some(ProtocolVersion::V_2026_07_28) + ); + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult( + DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + ServerCapabilities::default(), + ) + .with_server_info(Implementation::new("discover-server", "1.0.0")), + ), + second.id, + )) + .await + .expect("send discover response"); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ + serde_json::from_value(serde_json::json!("2099-01-01")).unwrap(), + ProtocolVersion::V_2026_07_28, + ], + }, + ) + .await + .expect("discover client should retry"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn discover_startup_retries_current_version_once_when_server_reports_it_supported() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(first) = + server.receive().await.expect("expected first discover") + else { + panic!("expected request"); + }; + server + .send(ServerJsonRpcMessage::error( + ErrorData::unsupported_protocol_version( + ProtocolVersion::V_2026_07_28, + &[ProtocolVersion::V_2026_07_28], + ), + Some(first.id), + )) + .await + .expect("send unsupported error"); + + let ClientJsonRpcMessage::Request(second) = + server.receive().await.expect("expected retry discover") + else { + panic!("expected request"); + }; + assert_eq!( + second.request.get_meta().protocol_version(), + Some(ProtocolVersion::V_2026_07_28) + ); + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult( + DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + ServerCapabilities::default(), + ) + .with_server_info(Implementation::new("discover-server", "1.0.0")), + ), + second.id, + )) + .await + .expect("send discover response"); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("discover client should retry once"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} diff --git a/crates/rmcp/tests/test_custom_headers.rs b/crates/rmcp/tests/test_custom_headers.rs index 736dce1..b01223a 100644 --- a/crates/rmcp/tests/test_custom_headers.rs +++ b/crates/rmcp/tests/test_custom_headers.rs @@ -380,10 +380,10 @@ async fn test_mcp_custom_headers_sent_to_server() -> anyhow::Result<()> { let mut headers_map = HashMap::new(); for (name, value) in headers.iter() { let name_str = name.as_str(); - if name_str.starts_with("x-") { - if let Ok(v) = value.to_str() { - headers_map.insert(name_str.to_string(), v.to_string()); - } + if name_str.starts_with("x-") + && let Ok(v) = value.to_str() + { + headers_map.insert(name_str.to_string(), v.to_string()); } } @@ -392,48 +392,48 @@ async fn test_mcp_custom_headers_sent_to_server() -> anyhow::Result<()> { stored.extend(headers_map); // Parse the MCP request - if let Ok(json_body) = serde_json::from_slice::(&body) { - if let Some(method) = json_body.get("method").and_then(|m| m.as_str()) { - if method == "initialize" { - state.initialize_called.notify_one(); - // Return a valid MCP initialize response with session header - let response = json!({ - "jsonrpc": "2.0", - "id": json_body.get("id"), - "result": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "serverInfo": { - "name": "test-server", - "version": "1.0.0" - } + if let Ok(json_body) = serde_json::from_slice::(&body) + && let Some(method) = json_body.get("method").and_then(|m| m.as_str()) + { + if method == "initialize" { + state.initialize_called.notify_one(); + // Return a valid MCP initialize response with session header + let response = json!({ + "jsonrpc": "2.0", + "id": json_body.get("id"), + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "serverInfo": { + "name": "test-server", + "version": "1.0.0" } - }); - return ( - StatusCode::OK, - [ - (http::header::CONTENT_TYPE, "application/json"), - ( - http::HeaderName::from_static("mcp-session-id"), - "test-session-123", - ), - ], - response.to_string(), - ); - } else if method == "notifications/initialized" { - // For initialized notification, return 202 Accepted - return ( - StatusCode::ACCEPTED, - [ - (http::header::CONTENT_TYPE, "application/json"), - ( - http::HeaderName::from_static("mcp-session-id"), - "test-session-123", - ), - ], - String::new(), - ); - } + } + }); + return ( + StatusCode::OK, + [ + (http::header::CONTENT_TYPE, "application/json"), + ( + http::HeaderName::from_static("mcp-session-id"), + "test-session-123", + ), + ], + response.to_string(), + ); + } else if method == "notifications/initialized" { + // For initialized notification, return 202 Accepted + return ( + StatusCode::ACCEPTED, + [ + (http::header::CONTENT_TYPE, "application/json"), + ( + http::HeaderName::from_static("mcp-session-id"), + "test-session-123", + ), + ], + String::new(), + ); } } @@ -719,7 +719,7 @@ async fn test_server_rejects_unsupported_protocol_version() { use http_body_util::Full; use rmcp::{ handler::server::ServerHandler, - model::{ServerCapabilities, ServerInfo}, + model::{ServerCapabilities, ServerConfig}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -730,8 +730,8 @@ async fn test_server_rejects_unsupported_protocol_version() { struct TestHandler; impl ServerHandler for TestHandler { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().build()) } } @@ -880,10 +880,11 @@ fn test_protocol_version_utilities() { assert!(ProtocolVersion::KNOWN_VERSIONS.contains(&ProtocolVersion::V_2026_07_28)); } -/// Integration test: Verify server validates only the Host header for DNS rebinding protection +/// Integration test: Verify Host validation remains enabled when Origin validation is +/// disabled by default #[tokio::test] #[cfg(all(feature = "transport-streamable-http-server", feature = "server",))] -async fn test_server_validates_host_header_for_dns_rebinding_protection() { +async fn test_server_validates_host_when_origin_validation_is_disabled_by_default() { use std::sync::Arc; use bytes::Bytes; @@ -891,7 +892,7 @@ async fn test_server_validates_host_header_for_dns_rebinding_protection() { use http_body_util::Full; use rmcp::{ handler::server::ServerHandler, - model::{ServerCapabilities, ServerInfo}, + model::{ServerCapabilities, ServerConfig}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -902,8 +903,8 @@ async fn test_server_validates_host_header_for_dns_rebinding_protection() { struct TestHandler; impl ServerHandler for TestHandler { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().build()) } } @@ -974,7 +975,7 @@ async fn test_server_validates_host_header_port_for_dns_rebinding_protection() { use http_body_util::Full; use rmcp::{ handler::server::ServerHandler, - model::{ServerCapabilities, ServerInfo}, + model::{ServerCapabilities, ServerConfig}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -985,8 +986,8 @@ async fn test_server_validates_host_header_port_for_dns_rebinding_protection() { struct TestHandler; impl ServerHandler for TestHandler { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().build()) } } @@ -1045,7 +1046,7 @@ async fn test_server_falls_back_to_uri_authority_when_host_header_missing() { use http_body_util::Full; use rmcp::{ handler::server::ServerHandler, - model::{ServerCapabilities, ServerInfo}, + model::{ServerCapabilities, ServerConfig}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -1056,8 +1057,8 @@ async fn test_server_falls_back_to_uri_authority_when_host_header_missing() { struct TestHandler; impl ServerHandler for TestHandler { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().build()) } } @@ -1127,11 +1128,11 @@ mod origin_validation { use std::sync::Arc; use bytes::Bytes; - use http::{Method, Request, header::CONTENT_TYPE}; + use http::{HeaderValue, Method, Request, header::CONTENT_TYPE}; use http_body_util::Full; use rmcp::{ handler::server::ServerHandler, - model::{ServerCapabilities, ServerInfo}, + model::{ServerCapabilities, ServerConfig}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -1142,17 +1143,25 @@ mod origin_validation { struct TestHandler; impl ServerHandler for TestHandler { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().build()) } } + fn service_with_config( + config: StreamableHttpServerConfig, + ) -> StreamableHttpService { + StreamableHttpService::new( + || Ok(TestHandler), + Arc::new(LocalSessionManager::default()), + config, + ) + } + fn service_with_allowed_origins( origins: &[&str], ) -> StreamableHttpService { - StreamableHttpService::new( - || Ok(TestHandler), - Arc::new(LocalSessionManager::default()), + service_with_config( StreamableHttpServerConfig::default().with_allowed_origins(origins.iter().copied()), ) } @@ -1199,6 +1208,80 @@ mod origin_validation { assert_eq!(response.status(), http::StatusCode::FORBIDDEN); } + #[tokio::test] + async fn malformed_origin_is_forbidden() { + let service = service_with_allowed_origins(&["http://localhost:8080"]); + let response = service.handle(init_request(Some("not-an-origin"))).await; + assert_eq!(response.status(), http::StatusCode::FORBIDDEN); + } + + #[tokio::test] + async fn non_utf8_origin_is_forbidden() { + let service = service_with_allowed_origins(&["http://localhost:8080"]); + let mut request = init_request(None); + request.headers_mut().insert( + http::header::ORIGIN, + HeaderValue::from_bytes(b"\xff").unwrap(), + ); + let response = service.handle(request).await; + assert_eq!(response.status(), http::StatusCode::FORBIDDEN); + } + + #[tokio::test] + async fn enabled_empty_allowlist_forbids_present_origin() { + let service = + service_with_config(StreamableHttpServerConfig::default().enforce_origin_validation()); + let response = service + .handle(init_request(Some("http://localhost:8080"))) + .await; + assert_eq!(response.status(), http::StatusCode::FORBIDDEN); + } + + #[tokio::test] + async fn enabled_empty_allowlist_allows_missing_origin() { + let service = + service_with_config(StreamableHttpServerConfig::default().enforce_origin_validation()); + let response = service.handle(init_request(None)).await; + assert_eq!(response.status(), http::StatusCode::OK); + } + + #[tokio::test] + async fn empty_allowed_origins_preserves_disabled_validation() { + let service = service_with_config( + StreamableHttpServerConfig::default().with_allowed_origins(std::iter::empty::<&str>()), + ); + let response = service + .handle(init_request(Some("http://attacker.example"))) + .await; + assert_eq!(response.status(), http::StatusCode::OK); + } + + #[tokio::test] + async fn nonempty_public_allowed_origins_field_enables_validation() { + let mut config = StreamableHttpServerConfig::default(); + config + .allowed_origins + .push("http://localhost:8080".to_string()); + let service = service_with_config(config); + let response = service + .handle(init_request(Some("http://attacker.example"))) + .await; + assert_eq!(response.status(), http::StatusCode::FORBIDDEN); + } + + #[tokio::test] + async fn explicitly_disabled_validation_allows_present_origin() { + let service = service_with_config( + StreamableHttpServerConfig::default() + .enforce_origin_validation() + .disable_allowed_origins(), + ); + let response = service + .handle(init_request(Some("http://attacker.example"))) + .await; + assert_eq!(response.status(), http::StatusCode::OK); + } + #[tokio::test] async fn missing_origin_passes_through() { let service = service_with_allowed_origins(&["http://localhost:8080"]); diff --git a/crates/rmcp/tests/test_deserialization.rs b/crates/rmcp/tests/test_deserialization.rs index 58e9a58..8570770 100644 --- a/crates/rmcp/tests/test_deserialization.rs +++ b/crates/rmcp/tests/test_deserialization.rs @@ -17,10 +17,8 @@ fn test_tool_list_result() { /// Regression tests for `#[serde(untagged)]` deserialization of `ServerResult`. /// /// `ServerResult` is an untagged enum, so serde tries each variant in declaration -/// order. `GetTaskPayloadResult` has a custom `Deserialize` impl that always fails -/// so it is skipped, and `CustomResult(Value)` acts as the catch-all. If variant -/// ordering changes or the custom impl is removed, these tests will catch the -/// regression. +/// order, with `CustomResult(Value)` acting as the catch-all. If variant ordering +/// changes, these tests will catch the regression. mod untagged_server_result { use rmcp::model::{CallToolResult, JsonRpcResponse, ServerJsonRpcMessage, ServerResult}; use serde_json::json; @@ -72,6 +70,81 @@ mod untagged_server_result { ); } + #[test] + fn input_required_result_with_meta_deserializes_to_correct_variant() { + let result = parse_result(wrap_response(json!({ + "resultType": "input_required", + "inputRequests": { + "username": { + "method": "elicitation/create", + "params": { + "message": "Please provide your username", + "requestedSchema": { + "type": "object", + "properties": { + "username": { "type": "string" } + }, + "required": ["username"] + } + } + } + }, + "requestState": "opaque-state", + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "test-server", + "version": "1.0.0" + } + } + }))); + + let ServerResult::InputRequiredResult(result) = result else { + panic!("expected InputRequiredResult, got {result:?}"); + }; + assert!( + result.input_requests.is_some_and(|requests| { + requests.len() == 1 && requests.contains_key("username") + }) + ); + assert_eq!(result.request_state.as_deref(), Some("opaque-state")); + assert_eq!( + result + .meta + .as_ref() + .and_then(|meta| meta.get("io.modelcontextprotocol/serverInfo")), + Some(&json!({ + "name": "test-server", + "version": "1.0.0" + })) + ); + } + + #[test] + fn call_tool_result_rejects_input_required_discriminator() { + assert!( + serde_json::from_value::(json!({ + "resultType": "input_required", + "requestState": "opaque-state", + "_meta": {} + })) + .is_err() + ); + } + + #[test] + fn invalid_input_required_result_falls_through_to_custom_result() { + let payload = json!({ + "resultType": "input_required", + "_meta": {} + }); + let result = parse_result(wrap_response(payload.clone())); + + let ServerResult::CustomResult(result) = result else { + panic!("expected CustomResult, got {result:?}"); + }; + assert_eq!(result.0, payload); + } + #[test] fn empty_object_deserializes_to_empty_result() { let result = parse_result(wrap_response(json!({}))); @@ -84,7 +157,7 @@ mod untagged_server_result { #[test] fn unknown_shape_falls_through_to_custom_result() { // A value that doesn't match any known result type should land in - // CustomResult, NOT GetTaskPayloadResult. + // CustomResult. let result = parse_result(wrap_response(json!({ "some_unknown_field": "some_value", "number": 42 @@ -96,10 +169,40 @@ mod untagged_server_result { } #[test] - fn arbitrary_json_value_does_not_deserialize_as_get_task_payload_result() { - // GetTaskPayloadResult wraps a bare Value, but its custom Deserialize - // always fails so serde skips it during untagged resolution. - // Any JSON value must fall through to CustomResult instead. + fn result_type_bearing_objects_do_not_match_task_ack() { + // TaskAckResult carries only `resultType` (+ optional `_meta`), so it + // must not greedily swallow arbitrary results that happen to include + // a `resultType` key inside the untagged ServerResult union. + let result = parse_result(wrap_response(json!({ + "resultType": "weird-custom", + "payload": { "a": 1 } + }))); + assert!( + matches!(result, ServerResult::CustomResult(_)), + "expected CustomResult, got {result:?}" + ); + + let result = parse_result(wrap_response(json!({ + "resultType": "complete", + "customField": 42 + }))); + assert!( + matches!(result, ServerResult::CustomResult(_)), + "expected CustomResult, got {result:?}" + ); + + // A bare complete ack (the actual tasks/update / tasks/cancel ack + // shape) still parses as TaskAckResult. + let result = parse_result(wrap_response(json!({ "resultType": "complete" }))); + assert!( + matches!(result, ServerResult::TaskAckResult(_)), + "expected TaskAckResult, got {result:?}" + ); + } + + #[test] + fn arbitrary_json_value_falls_through_to_custom_result() { + // Any bare JSON value must fall through to CustomResult. for value in [json!(42), json!("hello"), json!(null), json!([1, 2, 3])] { let result = parse_result(wrap_response(value.clone())); assert!( diff --git a/crates/rmcp/tests/test_discover_http_client_startup.rs b/crates/rmcp/tests/test_discover_http_client_startup.rs new file mode 100644 index 0000000..2d92df5 --- /dev/null +++ b/crates/rmcp/tests/test_discover_http_client_startup.rs @@ -0,0 +1,238 @@ +#![cfg(all( + not(feature = "local"), + feature = "client", + feature = "reqwest", + feature = "transport-streamable-http-server" +))] + +use std::{borrow::Cow, sync::Arc}; + +use axum::{ + Router, + body::{Body, Bytes}, + extract::State, + http::{Response, StatusCode}, + routing::post, +}; +use rmcp::{ + ClientLifecycleMode, ClientServiceExt, ServerHandler, + model::{ClientConfig, DiscoverResult, ErrorCode, ErrorData, ProtocolVersion}, + service::{MaybeSendFuture, RequestContext, RoleServer}, + transport::{ + StreamableHttpClientTransport, + streamable_http_client::StreamableHttpClientTransportConfig, + streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, + }, + }, +}; +use serde_json::json; +use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; + +#[derive(Clone, Default)] +struct DiscoverHttpServer; + +impl ServerHandler for DiscoverHttpServer { + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(&[ProtocolVersion::V_2026_07_28]) + } +} + +#[derive(Clone, Default)] +struct LegacyHttpServer; + +impl ServerHandler for LegacyHttpServer { + fn discover( + &self, + _context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + std::future::ready(Err(ErrorData::new( + ErrorCode::METHOD_NOT_FOUND, + "Method not found", + None, + ))) + } +} + +#[derive(Clone, Default)] +struct PlainTextLegacyHttpState { + methods: Arc>>, +} + +async fn plain_text_legacy_http_handler( + State(state): State, + body: Bytes, +) -> Response { + let request: serde_json::Value = serde_json::from_slice(&body).expect("valid JSON-RPC body"); + let method = request["method"] + .as_str() + .expect("request method") + .to_owned(); + state.methods.lock().await.push(method.clone()); + + if method == "server/discover" { + return Response::builder() + .status(StatusCode::UNPROCESSABLE_ENTITY) + .body(Body::from("Unexpected message, expect initialize request")) + .expect("build rejection response"); + } + + if method == "initialize" { + return Response::builder() + .status(StatusCode::OK) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "jsonrpc": "2.0", + "id": request["id"], + "result": { + "protocolVersion": "2025-11-25", + "capabilities": {}, + "serverInfo": {"name": "legacy", "version": "1.0"} + } + }) + .to_string(), + )) + .expect("build initialize response"); + } + + Response::builder() + .status(StatusCode::ACCEPTED) + .body(Body::empty()) + .expect("build notification response") +} + +#[tokio::test] +async fn discover_http_client_bootstraps_headers_without_initialize() { + let ct = CancellationToken::new(); + let service: StreamableHttpService = + StreamableHttpService::new( + || Ok(DiscoverHttpServer), + Default::default(), + StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_json_response(true) + .with_cancellation_token(ct.child_token()), + ); + let router = axum::Router::new().nest_service("/mcp", service); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let address = listener.local_addr().expect("listener address"); + let server = tokio::spawn({ + let ct = ct.clone(); + async move { + let _ = axum::serve(listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(format!("http://{address}/mcp")), + ); + let client = ClientConfig::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("discover HTTP client should start"); + client.list_tools(None).await.expect("list tools"); + client.cancel().await.expect("cancel client"); + + ct.cancel(); + server.await.expect("server task"); +} + +#[tokio::test] +async fn auto_http_client_falls_back_to_stateful_legacy_startup() { + let ct = CancellationToken::new(); + let service: StreamableHttpService = + StreamableHttpService::new( + || Ok(LegacyHttpServer), + Default::default(), + StreamableHttpServerConfig::default() + .with_json_response(true) + .with_cancellation_token(ct.child_token()), + ); + let router = axum::Router::new().nest_service("/mcp", service); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let address = listener.local_addr().expect("listener address"); + let server = tokio::spawn({ + let ct = ct.clone(); + async move { + let _ = axum::serve(listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(format!("http://{address}/mcp")), + ); + let client = ClientConfig::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await + .expect("auto HTTP client should fall back"); + client.list_tools(None).await.expect("list tools"); + client.cancel().await.expect("cancel client"); + + ct.cancel(); + server.await.expect("server task"); +} + +#[tokio::test] +async fn auto_http_client_falls_back_after_plain_text_4xx_rejection() { + let ct = CancellationToken::new(); + let state = PlainTextLegacyHttpState::default(); + let methods = state.methods.clone(); + let router = Router::new() + .route("/mcp", post(plain_text_legacy_http_handler)) + .with_state(state); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let address = listener.local_addr().expect("listener address"); + let server = tokio::spawn({ + let ct = ct.clone(); + async move { + let _ = axum::serve(listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(format!("http://{address}/mcp")), + ); + let client = ClientConfig::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await + .expect("auto HTTP client should fall back after a transport-level rejection"); + client.cancel().await.expect("cancel client"); + + assert_eq!( + methods.lock().await.as_slice(), + &["server/discover", "initialize", "notifications/initialized"] + ); + ct.cancel(); + server.await.expect("server task"); +} diff --git a/crates/rmcp/tests/test_elicitation.rs b/crates/rmcp/tests/test_elicitation.rs index b4a1633..953648c 100644 --- a/crates/rmcp/tests/test_elicitation.rs +++ b/crates/rmcp/tests/test_elicitation.rs @@ -90,7 +90,7 @@ async fn test_elicitation_request_param_serialization() { assert_eq!(msg1, msg2); assert_eq!(schema1, schema2); } - _ => panic!("Expected FormElicitationParam variant"), + _ => panic!("Expected FormElicitationParams variant"), } } @@ -125,7 +125,7 @@ async fn test_elicitation_result_serialization() { assert_eq!(deserialized.meta, None); // Test protocol-level metadata round-trips as _meta. - let meta_result = ElicitResult::new(ElicitationAction::Accept).with_meta(Meta(object!({ + let meta_result = ElicitResult::new(ElicitationAction::Accept).with_meta(MetaObject(object!({ "traceId": "elicitation-123" }))); @@ -139,7 +139,7 @@ async fn test_elicitation_result_serialization() { let deserialized: ElicitResult = serde_json::from_value(expected).unwrap(); assert_eq!( deserialized.meta, - Some(Meta(object!({ "traceId": "elicitation-123" }))) + Some(MetaObject(object!({ "traceId": "elicitation-123" }))) ); } @@ -179,7 +179,7 @@ async fn test_elicitation_json_rpc_protocol() { ElicitRequestParams::FormElicitationParams { message, .. } => { assert_eq!(message, "Do you want to continue?"); } - _ => panic!("Expected FormElicitationParam variant"), + _ => panic!("Expected FormElicitationParams variant"), } } @@ -496,7 +496,7 @@ async fn test_elicitation_structured_schemas() { ]) ); } - _ => panic!("Expected FormElicitationParam variant"), + _ => panic!("Expected FormElicitationParams variant"), } } @@ -738,7 +738,7 @@ async fn test_elicitation_multi_select_enum() { ) } } - _ => panic!("Expected FormElicitationParam variant"), + _ => panic!("Expected FormElicitationParams variant"), } } @@ -799,7 +799,7 @@ async fn test_elicitation_single_select_enum() { ) } } - _ => panic!("Expected FormElicitationParam variant"), + _ => panic!("Expected FormElicitationParams variant"), } } @@ -1056,10 +1056,8 @@ async fn test_client_capabilities_with_elicitation() { assert!(capabilities_without.elicitation.is_none()); } -/// Test InitializeRequestParam with elicitation capability #[tokio::test] async fn test_initialize_request_with_elicitation() { - // Test InitializeRequestParam with elicitation capability let init_param = InitializeRequestParams::new( ClientCapabilities::builder() .enable_elicitation_with( @@ -1829,7 +1827,7 @@ async fn test_url_elicitation_request_param_serialization() { assert_eq!(url, "https://example.com/verify"); assert_eq!(elicitation_id, "elicit-123"); } - _ => panic!("Expected UrlElicitationParam variant"), + _ => panic!("Expected UrlElicitationParams variant"), } } @@ -1878,35 +1876,10 @@ async fn test_url_elicitation_json_rpc_protocol() { assert_eq!(url, "https://auth.example.com/authorize/abc123"); assert_eq!(elicitation_id, "auth-request-456"); } - _ => panic!("Expected UrlElicitationParam variant"), + _ => panic!("Expected UrlElicitationParams variant"), } } -/// Test ElicitationCompleteNotification serialization/deserialization -#[tokio::test] -async fn test_elicitation_completion_notification() { - let notification_params = ElicitationResponseNotificationParam::new("elicit-789"); - - // Test serialization - let json = serde_json::to_value(¬ification_params).unwrap(); - let expected = json!({ - "elicitationId": "elicit-789" - }); - assert_eq!(json, expected); - - // Test deserialization - let deserialized: ElicitationResponseNotificationParam = - serde_json::from_value(expected).unwrap(); - assert_eq!(deserialized.elicitation_id, "elicit-789"); - - // Test complete notification structure - let notification = ElicitationCompleteNotification::new(notification_params); - - let json = serde_json::to_value(¬ification).unwrap(); - assert_eq!(json["method"], "notifications/elicitation/complete"); - assert_eq!(json["params"]["elicitationId"], "elicit-789"); -} - /// Test UrlElicitationCapability structure and serialization #[tokio::test] async fn test_url_elicitation_capability() { @@ -1952,7 +1925,6 @@ async fn test_url_elicitation_capability() { /// Test backward compatibility: ElicitRequestParams without mode tag #[tokio::test] async fn test_elicitation_backward_compatibility_no_mode() { - // JSON without "mode" field should deserialize as FormElicitationParam let json_without_mode = json!({ "message": "Please enter your details", "requestedSchema": { @@ -1978,7 +1950,7 @@ async fn test_elicitation_backward_compatibility_no_mode() { assert_eq!(requested_schema.properties.len(), 1); assert!(requested_schema.properties.contains_key("name")); } - _ => panic!("Expected FormElicitationParam for backward compatibility"), + _ => panic!("Expected FormElicitationParams for backward compatibility"), } } @@ -2020,39 +1992,6 @@ async fn test_elicitation_both_modes() { assert!(url_json.get("requestedSchema").is_none()); } -/// Test URL_ELICITATION_REQUIRED error code -#[tokio::test] -async fn test_url_elicitation_required_error_code() { - // Test the error code constant - assert_eq!(ErrorCode::URL_ELICITATION_REQUIRED.0, -32042); - - // Test creating error data with URL_ELICITATION_REQUIRED - let error_data = ErrorData::url_elicitation_required( - "URL elicitation is required for this operation", - Some(json!({ - "url": "https://example.com/complete", - "elicitationId": "elicit-999" - })), - ); - - assert_eq!(error_data.code, ErrorCode::URL_ELICITATION_REQUIRED); - assert_eq!( - error_data.message, - "URL elicitation is required for this operation" - ); - assert!(error_data.data.is_some()); - - // Test serialization - let json = serde_json::to_value(&error_data).unwrap(); - assert_eq!(json["code"], -32042); - assert_eq!( - json["message"], - "URL elicitation is required for this operation" - ); - assert_eq!(json["data"]["url"], "https://example.com/complete"); - assert_eq!(json["data"]["elicitationId"], "elicit-999"); -} - /// Test ClientCapabilities with different elicitation mode combinations #[tokio::test] async fn test_client_capabilities_elicitation_modes() { @@ -2102,32 +2041,6 @@ async fn test_client_capabilities_elicitation_modes() { assert!(json["elicitation"]["url"].is_object()); } -/// Test ElicitationCompleteNotification in ServerNotification enum -#[tokio::test] -async fn test_elicitation_completion_in_server_notification() { - let notification_param = ElicitationResponseNotificationParam::new("notify-123"); - - let completion_notification = ElicitationCompleteNotification::new(notification_param.clone()); - - // Test that it's part of ServerNotification - let server_notification = - ServerNotification::ElicitationCompleteNotification(completion_notification); - - // Test serialization - let json = serde_json::to_value(&server_notification).unwrap(); - assert_eq!(json["method"], "notifications/elicitation/complete"); - assert_eq!(json["params"]["elicitationId"], "notify-123"); - - // Test deserialization - let deserialized: ServerNotification = serde_json::from_value(json).unwrap(); - match deserialized { - ServerNotification::ElicitationCompleteNotification(notif) => { - assert_eq!(notif.params.elicitation_id, "notify-123"); - } - _ => panic!("Expected ElicitationCompleteNotification variant"), - } -} - /// Test ElicitationAction with URL elicitation workflow #[tokio::test] async fn test_url_elicitation_action_workflow() { @@ -2161,10 +2074,4 @@ async fn test_elicitation_method_constants() { ElicitationResponseNotificationMethod::VALUE, "notifications/elicitation/response" ); - - // Test new completion notification method - assert_eq!( - ElicitationCompletionNotificationMethod::VALUE, - "notifications/elicitation/complete" - ); } diff --git a/crates/rmcp/tests/test_embedded_resource_meta.rs b/crates/rmcp/tests/test_embedded_resource_meta.rs index 167108e..5854f3f 100644 --- a/crates/rmcp/tests/test_embedded_resource_meta.rs +++ b/crates/rmcp/tests/test_embedded_resource_meta.rs @@ -1,12 +1,12 @@ -use rmcp::model::{ContentBlock, EmbeddedResource, Meta, ResourceContents}; +use rmcp::model::{ContentBlock, EmbeddedResource, MetaObject, ResourceContents}; use serde_json::json; #[test] fn serialize_embedded_text_resource_with_meta() { - let mut resource_content_meta = Meta::new(); + let mut resource_content_meta = MetaObject::new(); resource_content_meta.insert("inner".to_string(), json!(2)); - let mut resource_meta = Meta::new(); + let mut resource_meta = MetaObject::new(); resource_meta.insert("top".to_string(), json!(1)); let content = ContentBlock::Resource( @@ -90,10 +90,10 @@ fn deserialize_embedded_text_resource_with_meta() { #[test] fn serialize_embedded_blob_resource_with_meta() { - let mut resource_content_meta = Meta::new(); + let mut resource_content_meta = MetaObject::new(); resource_content_meta.insert("blob_inner".to_string(), json!(true)); - let mut resource_meta = Meta::new(); + let mut resource_meta = MetaObject::new(); resource_meta.insert("blob_top".to_string(), json!("t")); let content = ContentBlock::Resource( diff --git a/crates/rmcp/tests/test_handler_cache_hints.rs b/crates/rmcp/tests/test_handler_cache_hints.rs new file mode 100644 index 0000000..ed3cf67 --- /dev/null +++ b/crates/rmcp/tests/test_handler_cache_hints.rs @@ -0,0 +1,122 @@ +#![cfg(not(feature = "local"))] +#![cfg(feature = "client")] + +use rmcp::{ + ClientHandler, RoleClient, RoleServer, ServerHandler, + handler::server::router::{prompt::PromptRouter, tool::ToolRouter}, + model::{ + CacheScope, ClientConfig, InitializeResult, ListPromptsResult, ListToolsResult, + ProtocolVersion, + }, + prompt_handler, + service::serve_directly, + tool_handler, +}; + +#[derive(Debug, Clone)] +struct CacheHintServer { + tool_router: ToolRouter, + prompt_router: PromptRouter, +} + +impl CacheHintServer { + fn new() -> Self { + Self { + tool_router: ToolRouter::new(), + prompt_router: PromptRouter::new(), + } + } +} + +#[tool_handler(router = self.tool_router)] +#[prompt_handler(router = self.prompt_router)] +impl ServerHandler for CacheHintServer {} + +#[derive(Debug, Clone)] +struct VersionedClient { + protocol_version: ProtocolVersion, +} + +impl ClientHandler for VersionedClient { + fn get_info(&self) -> ClientConfig { + let mut info = ClientConfig::default(); + info.protocol_version = self.protocol_version.clone(); + info + } +} + +/// Wires the pair up directly on `protocol_version`. `2026-07-28` removed the +/// `initialize` handshake, so a peer on that revision is reached the way the +/// discover lifecycle leaves one: with the version already agreed. +async fn list_results(protocol_version: ProtocolVersion) -> (ListToolsResult, ListPromptsResult) { + let (server_transport, client_transport) = tokio::io::duplex(4096); + + let client_handler = VersionedClient { + protocol_version: protocol_version.clone(), + }; + let mut server_peer_info = InitializeResult::default(); + server_peer_info.protocol_version = protocol_version; + + let server = serve_directly::( + CacheHintServer::new(), + server_transport, + Some(client_handler.get_info()), + ); + let server_handle = tokio::spawn(async move { + server.waiting().await?; + anyhow::Ok(()) + }); + + let client = serve_directly::( + client_handler, + client_transport, + Some(server_peer_info.into()), + ); + let tools = client + .list_tools(None) + .await + .expect("tools/list should succeed"); + let prompts = client + .list_prompts(None) + .await + .expect("prompts/list should succeed"); + + client.cancel().await.expect("client should cancel"); + server_handle.await.expect("server task").expect("server"); + (tools, prompts) +} + +#[tokio::test] +async fn handler_macros_should_emit_required_cache_hints_for_2026_07_28() { + let (tools, prompts) = list_results(ProtocolVersion::V_2026_07_28).await; + + assert_eq!( + ( + tools.ttl_ms, + tools.cache_scope, + prompts.ttl_ms, + prompts.cache_scope, + ), + ( + Some(0), + Some(CacheScope::Public), + Some(0), + Some(CacheScope::Public), + ) + ); +} + +#[tokio::test] +async fn handler_macros_should_omit_cache_hints_for_legacy_versions() { + let (tools, prompts) = list_results(ProtocolVersion::V_2025_11_25).await; + + assert_eq!( + ( + tools.ttl_ms, + tools.cache_scope, + prompts.ttl_ms, + prompts.cache_scope, + ), + (None, None, None, None) + ); +} diff --git a/crates/rmcp/tests/test_inflight_response_drain.rs b/crates/rmcp/tests/test_inflight_response_drain.rs index c75e17c..fc5908f 100644 --- a/crates/rmcp/tests/test_inflight_response_drain.rs +++ b/crates/rmcp/tests/test_inflight_response_drain.rs @@ -14,7 +14,7 @@ use std::{ use rmcp::{ ServerHandler, ServiceExt, handler::server::{router::tool::ToolRouter, wrapper::Parameters}, - model::{CallToolRequestParams, ClientInfo, ServerCapabilities, ServerInfo}, + model::{CallToolRequestParams, ClientConfig, ServerCapabilities, ServerConfig}, service::QuitReason, tool, tool_handler, tool_router, }; @@ -55,8 +55,8 @@ impl SlowToolServer { #[tool_handler] impl ServerHandler for SlowToolServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) } } @@ -64,8 +64,8 @@ impl ServerHandler for SlowToolServer { struct DummyClientHandler; impl rmcp::ClientHandler for DummyClientHandler { - fn get_info(&self) -> ClientInfo { - ClientInfo::default() + fn get_info(&self) -> ClientConfig { + ClientConfig::default() } } diff --git a/crates/rmcp/tests/test_json_schema_detection.rs b/crates/rmcp/tests/test_json_schema_detection.rs index 5d982cd..e13125c 100644 --- a/crates/rmcp/tests/test_json_schema_detection.rs +++ b/crates/rmcp/tests/test_json_schema_detection.rs @@ -60,6 +60,28 @@ impl TestServer { pub async fn explicit_schema(&self) -> Result { Ok("test".to_string()) } + + /// Tool that returns Json> - array output schema + #[tool(name = "with-json-array")] + pub async fn with_json_array(&self) -> Result>, String> { + Ok(Json(vec![TestData { + value: "test".to_string(), + }])) + } + + /// Tool that returns Result>, ErrorData> - array output schema + #[tool(name = "result-with-json-array")] + pub async fn result_with_json_array(&self) -> Result>, rmcp::ErrorData> { + Ok(Json(vec![TestData { + value: "test".to_string(), + }])) + } + + /// Tool that returns Json - string output schema + #[tool(name = "with-json-string")] + pub async fn with_json_string(&self) -> Result, String> { + Ok(Json("test".to_string())) + } } #[tokio::test] @@ -113,3 +135,54 @@ async fn test_explicit_schema_override() { "Explicit output_schema attribute should work" ); } + +#[tokio::test] +async fn test_json_array_type_generates_schema() { + let server = TestServer::new(); + let tools = server.tool_router.list_all(); + + let array_tool = tools.iter().find(|t| t.name == "with-json-array").unwrap(); + assert!( + array_tool.output_schema.is_some(), + "Json> return type should generate output schema" + ); + let schema = array_tool.output_schema.as_ref().unwrap(); + assert_eq!( + schema.get("type").and_then(|v| v.as_str()), + Some("array"), + "Json> should produce an array schema" + ); +} + +#[tokio::test] +async fn test_result_with_json_array_generates_schema() { + let server = TestServer::new(); + let tools = server.tool_router.list_all(); + + let result_array_tool = tools + .iter() + .find(|t| t.name == "result-with-json-array") + .unwrap(); + assert!( + result_array_tool.output_schema.is_some(), + "Result>, ErrorData> return type should generate output schema" + ); +} + +#[tokio::test] +async fn test_json_string_type_generates_schema() { + let server = TestServer::new(); + let tools = server.tool_router.list_all(); + + let string_tool = tools.iter().find(|t| t.name == "with-json-string").unwrap(); + assert!( + string_tool.output_schema.is_some(), + "Json return type should generate output schema" + ); + let schema = string_tool.output_schema.as_ref().unwrap(); + assert_eq!( + schema.get("type").and_then(|v| v.as_str()), + Some("string"), + "Json should produce a string schema" + ); +} diff --git a/crates/rmcp/tests/test_list_tools_result/list_tools_result.json b/crates/rmcp/tests/test_list_tools_result/list_tools_result.json index 15325e8..c6616b5 100644 --- a/crates/rmcp/tests/test_list_tools_result/list_tools_result.json +++ b/crates/rmcp/tests/test_list_tools_result/list_tools_result.json @@ -1,5 +1,6 @@ { "result": { + "resultType": "complete", "tools": [ { "name": "add", diff --git a/crates/rmcp/tests/test_live_oauth_refresh.rs b/crates/rmcp/tests/test_live_oauth_refresh.rs new file mode 100644 index 0000000..543abe5 --- /dev/null +++ b/crates/rmcp/tests/test_live_oauth_refresh.rs @@ -0,0 +1,233 @@ +//! Live authorization-server checks for the refresh path. +//! +//! These are `#[ignore]`d, so they never run in CI. Provision the authorization +//! server with `./scripts/keycloak-oauth-fixture.sh`, then run them with +//! `cargo test -p rmcp --all-features --test test_live_oauth_refresh -- --ignored`. +//! `KC_BASE` overrides the server location for both the script and these tests. +#![cfg(feature = "auth")] + +use std::sync::Arc; + +use rmcp::transport::auth::{ + AuthError, AuthorizationManager, AuthorizationMetadata, CredentialRefreshGuard, + CredentialStore, InMemoryCredentialStore, OAuthClientConfig, OAuthTokenResponse, + StoredCredentials, +}; +use tokio::sync::Mutex; + +const REALM: &str = "rmcp"; +const CLIENT_ID: &str = "rmcp-client"; +const CLIENT_SECRET: &str = "rmcp-secret"; + +fn kc_base() -> String { + std::env::var("KC_BASE").unwrap_or_else(|_| "http://localhost:8081".to_string()) +} + +fn token_endpoint() -> String { + format!("{}/realms/{REALM}/protocol/openid-connect/token", kc_base()) +} + +/// Ask Keycloak for a genuine token pair through the direct access grant, so the +/// stored credentials hold a refresh token the server will actually honor. +async fn issue_real_credentials() -> OAuthTokenResponse { + let form = format!( + "client_id={CLIENT_ID}&client_secret={CLIENT_SECRET}\ + &username=alice&password=alice-pw&grant_type=password&scope=openid+profile" + ); + let body = reqwest::Client::new() + .post(token_endpoint()) + .header("content-type", "application/x-www-form-urlencoded") + .body(form) + .send() + .await + .expect("keycloak unreachable") + .text() + .await + .unwrap(); + serde_json::from_str(&body).unwrap_or_else(|e| panic!("unexpected token response {body}: {e}")) +} + +fn metadata() -> AuthorizationMetadata { + // `AuthorizationMetadata` is `#[non_exhaustive]`, so fill it field by field. + let mut metadata = AuthorizationMetadata::default(); + metadata.authorization_endpoint = + format!("{}/realms/{REALM}/protocol/openid-connect/auth", kc_base()); + metadata.token_endpoint = token_endpoint(); + metadata +} + +fn client_config() -> OAuthClientConfig { + let mut config = OAuthClientConfig::new(CLIENT_ID, "http://localhost/callback"); + config.client_secret = Some(CLIENT_SECRET.to_string()); + config +} + +async fn manager_with_store(store: S) -> AuthorizationManager { + let mut manager = AuthorizationManager::new(kc_base()).await.unwrap(); + manager.set_metadata(metadata()); + manager.configure_client(client_config()).unwrap(); + manager.set_credential_store(store); + manager +} + +fn stored(client_id: &str, token: OAuthTokenResponse) -> StoredCredentials { + StoredCredentials::new( + client_id.to_string(), + Some(token), + vec!["openid".into(), "profile".into()], + None, + ) +} + +/// A store that serializes refreshes the way a shared on-disk store would. +#[derive(Clone, Default)] +struct GuardedStore { + inner: InMemoryCredentialStore, + lock: Arc>, +} + +#[async_trait::async_trait] +impl CredentialStore for GuardedStore { + async fn load(&self) -> Result, AuthError> { + self.inner.load().await + } + + async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> { + self.inner.save(credentials).await + } + + async fn clear(&self) -> Result<(), AuthError> { + self.inner.clear().await + } + + async fn acquire_refresh_guard(&self) -> Result, AuthError> { + Ok(Some(CredentialRefreshGuard::new( + self.lock.clone().lock_owned().await, + ))) + } +} + +#[tokio::test] +#[ignore = "requires a live Keycloak"] +async fn live_refresh_rotates_the_stored_token() { + use oauth2::TokenResponse; + + let issued = issue_real_credentials().await; + let original_refresh = issued.refresh_token().unwrap().secret().clone(); + let store = InMemoryCredentialStore::new(); + store.save(stored(CLIENT_ID, issued)).await.unwrap(); + let manager = manager_with_store(store.clone()).await; + + let refreshed = manager.refresh_token().await.expect("live refresh failed"); + + let saved = store.load().await.unwrap().unwrap(); + let saved_token = saved.token_response.unwrap(); + assert_ne!( + saved_token.refresh_token().unwrap().secret(), + &original_refresh, + "keycloak rotates refresh tokens, so the store must hold the new one" + ); + assert_eq!( + saved_token.access_token().secret(), + refreshed.access_token().secret(), + "the saved credentials must match what the caller received" + ); + assert!( + saved.granted_scopes.contains(&"openid".to_string()), + "granted scopes should come from the provider response, got: {:?}", + saved.granted_scopes + ); +} + +#[tokio::test] +#[ignore = "requires a live Keycloak"] +async fn live_refresh_rejects_credentials_for_another_client() { + use oauth2::TokenResponse; + + let issued = issue_real_credentials().await; + let untouched_refresh = issued.refresh_token().unwrap().secret().clone(); + let store = InMemoryCredentialStore::new(); + store + .save(stored("some-other-client", issued)) + .await + .unwrap(); + let manager = manager_with_store(store.clone()).await; + + let error = manager.refresh_token().await.unwrap_err(); + + assert!( + matches!(error, AuthError::AuthorizationRequired), + "a client mismatch must require reauthorization, got: {error:?}" + ); + let saved = store.load().await.unwrap().unwrap(); + assert_eq!( + saved + .token_response + .unwrap() + .refresh_token() + .unwrap() + .secret(), + &untouched_refresh, + "a rejected refresh must leave the stored token untouched" + ); +} + +#[tokio::test] +#[ignore = "requires a live Keycloak"] +async fn live_concurrent_refreshes_survive_refresh_token_rotation() { + use oauth2::TokenResponse; + + let issued = issue_real_credentials().await; + let store = GuardedStore::default(); + store.save(stored(CLIENT_ID, issued)).await.unwrap(); + let first = manager_with_store(store.clone()).await; + let second = manager_with_store(store.clone()).await; + + // Without the guard the second caller would reuse the refresh token the first + // one already consumed, and Keycloak would answer `invalid_grant`. + let (a, b) = tokio::join!( + tokio::spawn(async move { first.refresh_token().await }), + tokio::spawn(async move { second.refresh_token().await }) + ); + + let a = a.unwrap().expect("first concurrent refresh failed"); + let b = b.unwrap().expect("second concurrent refresh failed"); + assert_ne!( + a.access_token().secret(), + b.access_token().secret(), + "each caller performs its own exchange, so the tokens must differ" + ); +} + +/// Deterministic proof that this realm really does invalidate a rotated refresh +/// token, which is what makes the coordination above load-bearing. +#[tokio::test] +#[ignore = "requires a live Keycloak with revokeRefreshToken enabled"] +async fn live_reusing_a_rotated_refresh_token_is_rejected() { + let issued = issue_real_credentials().await; + + let first_store = InMemoryCredentialStore::new(); + first_store + .save(stored(CLIENT_ID, issued.clone())) + .await + .unwrap(); + manager_with_store(first_store) + .await + .refresh_token() + .await + .expect("the first refresh should succeed"); + + // A second caller that never saw the rotation still holds the consumed token. + let stale_store = InMemoryCredentialStore::new(); + stale_store.save(stored(CLIENT_ID, issued)).await.unwrap(); + let error = manager_with_store(stale_store) + .await + .refresh_token() + .await + .unwrap_err(); + + assert!( + matches!(error, AuthError::TokenRefreshRejected(_)), + "reusing a rotated refresh token must be rejected, got: {error:?}" + ); +} diff --git a/crates/rmcp/tests/test_local_client_handler.rs b/crates/rmcp/tests/test_local_client_handler.rs new file mode 100644 index 0000000..e85d3d8 --- /dev/null +++ b/crates/rmcp/tests/test_local_client_handler.rs @@ -0,0 +1,16 @@ +use std::rc::Rc; + +use rmcp::ClientHandler; + +struct LocalClientHandler { + _state: Rc<()>, +} + +impl ClientHandler for LocalClientHandler {} + +#[test] +fn client_handler_accepts_non_send_sync_state_with_local_feature() { + fn assert_client_handler() {} + + assert_client_handler::(); +} diff --git a/crates/rmcp/tests/test_message_schema.rs b/crates/rmcp/tests/test_message_schema.rs index f05c263..027a2cc 100644 --- a/crates/rmcp/tests/test_message_schema.rs +++ b/crates/rmcp/tests/test_message_schema.rs @@ -61,12 +61,73 @@ mod tests { ); } + /// The three metadata definitions must expose the MCP 2026-07-28 + /// vocabulary: `MetaObject` is an open map, `RequestMetaObject` reserves + /// `progressToken` plus the SEP-2575 keys, and `NotificationMetaObject` + /// reserves `io.modelcontextprotocol/subscriptionId`. Version-specific + /// required keys stay optional because rmcp generates one schema shared + /// by every supported protocol version; current-version validation is + /// a runtime concern (`RequestMetaObject::missing_required_keys`). + #[test] + fn test_metadata_definitions_match_2026_07_28_schema() { + let settings = SchemaSettings::draft07(); + let schema = settings + .into_generator() + .into_root_schema_for::(); + let schema = serde_json::to_value(&schema).expect("Failed to serialize schema"); + let definitions = &schema["definitions"]; + + assert_eq!( + definitions["MetaObject"], + serde_json::json!({ + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", + "type": "object", + "additionalProperties": true, + }) + ); + + assert_eq!( + definitions["RequestMetaObject"], + serde_json::json!({ + "description": "Metadata reserved by MCP on requests. Extension keys are also allowed.", + "type": "object", + "properties": { + "progressToken": { "$ref": "#/definitions/ProgressToken" }, + "io.modelcontextprotocol/protocolVersion": { "type": "string" }, + "io.modelcontextprotocol/clientInfo": { "$ref": "#/definitions/Implementation" }, + "io.modelcontextprotocol/clientCapabilities": { "$ref": "#/definitions/ClientCapabilities" }, + "io.modelcontextprotocol/logLevel": { "$ref": "#/definitions/LoggingLevel" }, + }, + "additionalProperties": true, + }) + ); + + assert_eq!( + definitions["NotificationMetaObject"], + serde_json::json!({ + "description": "Metadata reserved by MCP on notifications. Extension keys are also allowed.", + "type": "object", + "properties": { + "io.modelcontextprotocol/subscriptionId": { "$ref": "#/definitions/NumberOrString" }, + }, + "additionalProperties": true, + }) + ); + } + #[test] fn test_server_json_rpc_message_schema() { let settings = SchemaSettings::draft07(); let schema = settings .into_generator() .into_root_schema_for::(); + let schema_value = + serde_json::to_value(&schema).expect("Failed to serialize server schema"); + let discover_result = &schema_value["definitions"]["DiscoverResult"]; + assert!( + discover_result["properties"].get("serverInfo").is_none(), + "DiscoverResult serverInfo belongs in namespaced _meta" + ); let schema_str = serde_json::to_string_pretty(&schema).expect("Failed to serialize schema"); compare_schemas( diff --git a/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema.json b/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema.json index 3bfdb7c..760d9e1 100644 --- a/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema.json +++ b/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema.json @@ -54,8 +54,7 @@ "type": [ "string", "null" - ], - "format": "date-time" + ] }, "priority": { "type": [ @@ -87,11 +86,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -124,19 +126,30 @@ "const": "tools/call" }, "CallToolRequestParams": { - "description": "Parameters for calling a tool provided by an MCP server.\n\nContains the tool name and optional arguments needed to execute\nthe tool operation.\n\nThis implements `TaskAugmentedRequestParamsMeta` as tool calls can be\nlong-running and may benefit from task-based execution.", + "description": "Parameters for calling a tool provided by an MCP server.\n\nContains the tool name and optional arguments needed to execute\nthe tool operation.", "type": "object", "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] + }, + "arguments": { + "description": "Arguments to pass to the tool (must match the tool's input schema)", "type": [ "object", "null" ], "additionalProperties": true }, - "arguments": { - "description": "Arguments to pass to the tool (must match the tool's input schema)", + "inputResponses": { + "description": "Client responses to server-initiated input requests from a previous\n[`InputRequiredResult`]. Present only when retrying after an incomplete result.", "type": [ "object", "null" @@ -147,15 +160,11 @@ "description": "The name of the tool to call", "type": "string" }, - "task": { - "description": "Task metadata for async task management (SEP-1319)", - "anyOf": [ - { - "$ref": "#/definitions/TaskMetadata" - }, - { - "type": "null" - } + "requestState": { + "description": "Opaque request state echoed back from a previous [`InputRequiredResult`].\nClients MUST return this value exactly as received.", + "type": [ + "string", + "null" ] } }, @@ -173,11 +182,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "taskId": { "type": "string" @@ -196,11 +208,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "reason": { "type": [ @@ -278,16 +293,6 @@ "type": "null" } ] - }, - "tasks": { - "anyOf": [ - { - "$ref": "#/definitions/TasksCapability" - }, - { - "type": "null" - } - ] } } }, @@ -320,11 +325,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "argument": { "$ref": "#/definitions/ArgumentInfo" @@ -460,11 +468,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "content": { "description": "The actual content of the message (text, image, audio, tool use, or tool result)", @@ -530,17 +541,36 @@ "CustomResult": { "description": "A catch-all response either side can use for custom requests." }, + "DiscoverRequestMethod": { + "type": "string", + "format": "const", + "const": "server/discover" + }, + "DiscoverRequestParams": { + "type": "object", + "properties": { + "_meta": { + "$ref": "#/definitions/RequestMetaObject" + } + }, + "required": [ + "_meta" + ] + }, "ElicitResult": { "description": "The result returned by a client in response to an elicitation request.\n\nContains the user's decision (accept/decline/cancel) and optionally their input data\nif they chose to accept the request.", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this result.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "action": { "description": "The user's decision on how to handle the elicitation request", @@ -606,29 +636,20 @@ } } }, - "ElicitationTaskCapability": { - "type": "object", - "properties": { - "create": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "EmbeddedResource": { "description": "Embedded resource content (spec `EmbeddedResource`).", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -713,11 +734,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "arguments": { "type": [ @@ -726,8 +750,23 @@ ], "additionalProperties": true }, + "inputResponses": { + "description": "Client responses to server-initiated input requests from a previous\n[`InputRequiredResult`].", + "type": [ + "object", + "null" + ], + "additionalProperties": true + }, "name": { "type": "string" + }, + "requestState": { + "description": "Opaque request state echoed back from a previous [`InputRequiredResult`].", + "type": [ + "string", + "null" + ] } }, "required": [ @@ -743,38 +782,17 @@ "type": "object", "properties": { "_meta": { - "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "taskId": { - "type": "string" - } - }, - "required": [ - "taskId" - ] - }, - "GetTaskPayloadMethod": { - "type": "string", - "format": "const", - "const": "tasks/result" - }, - "GetTaskPayloadParams": { - "type": "object", - "properties": { - "_meta": { - "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "taskId": { + "description": "Identifier of the task to query.", "type": "string" } }, @@ -844,11 +862,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -923,11 +944,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "capabilities": { "description": "The capabilities this client supports (sampling, roots, etc.)", @@ -1015,9 +1039,6 @@ { "$ref": "#/definitions/NotificationNoParam2" }, - { - "$ref": "#/definitions/Notification3" - }, { "$ref": "#/definitions/CustomNotification" } @@ -1052,6 +1073,9 @@ { "$ref": "#/definitions/Request4" }, + { + "$ref": "#/definitions/Request5" + }, { "$ref": "#/definitions/RequestOptionalParam" }, @@ -1061,9 +1085,6 @@ { "$ref": "#/definitions/RequestOptionalParam3" }, - { - "$ref": "#/definitions/Request5" - }, { "$ref": "#/definitions/Request6" }, @@ -1073,21 +1094,24 @@ { "$ref": "#/definitions/Request8" }, - { - "$ref": "#/definitions/RequestOptionalParam4" - }, { "$ref": "#/definitions/Request9" }, - { - "$ref": "#/definitions/RequestOptionalParam5" - }, { "$ref": "#/definitions/Request10" }, + { + "$ref": "#/definitions/RequestOptionalParam4" + }, { "$ref": "#/definitions/Request11" }, + { + "$ref": "#/definitions/Request12" + }, + { + "$ref": "#/definitions/Request13" + }, { "$ref": "#/definitions/CustomRequest" } @@ -1140,11 +1164,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "roots": { "type": "array", @@ -1158,11 +1185,6 @@ "roots" ] }, - "ListTasksMethod": { - "type": "string", - "format": "const", - "const": "tasks/list" - }, "ListToolsRequestMethod": { "type": "string", "format": "const", @@ -1183,6 +1205,11 @@ "emergency" ] }, + "MetaObject": { + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", + "type": "object", + "additionalProperties": true + }, "Notification": { "type": "object", "properties": { @@ -1213,20 +1240,15 @@ "params" ] }, - "Notification3": { + "NotificationMetaObject": { + "description": "Metadata reserved by MCP on notifications. Extension keys are also allowed.", "type": "object", "properties": { - "method": { - "$ref": "#/definitions/TaskStatusNotificationMethod" - }, - "params": { - "$ref": "#/definitions/TaskStatusNotificationParam" + "io.modelcontextprotocol/subscriptionId": { + "$ref": "#/definitions/NumberOrString" } }, - "required": [ - "method", - "params" - ] + "additionalProperties": true }, "NotificationNoParam": { "type": "object", @@ -1265,11 +1287,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "cursor": { "type": [ @@ -1293,11 +1318,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "message": { "description": "An optional message describing the current progress.", @@ -1368,12 +1396,30 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] + }, + "inputResponses": { + "description": "Client responses to server-initiated input requests from a previous\n[`InputRequiredResult`].", "type": [ "object", "null" ], "additionalProperties": true }, + "requestState": { + "description": "Opaque request state echoed back from a previous [`InputRequiredResult`].", + "type": [ + "string", + "null" + ] + }, "uri": { "description": "The URI of the resource to read", "type": "string" @@ -1442,10 +1488,10 @@ "type": "object", "properties": { "method": { - "$ref": "#/definitions/GetTaskPayloadMethod" + "$ref": "#/definitions/CallToolRequestMethod" }, "params": { - "$ref": "#/definitions/GetTaskPayloadParams" + "$ref": "#/definitions/CallToolRequestParams" } }, "required": [ @@ -1454,6 +1500,38 @@ ] }, "Request11": { + "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", + "type": "object", + "properties": { + "method": { + "$ref": "#/definitions/GetTaskMethod" + }, + "params": { + "$ref": "#/definitions/GetTaskParams" + } + }, + "required": [ + "method", + "params" + ] + }, + "Request12": { + "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", + "type": "object", + "properties": { + "method": { + "$ref": "#/definitions/UpdateTaskMethod" + }, + "params": { + "$ref": "#/definitions/UpdateTaskParams" + } + }, + "required": [ + "method", + "params" + ] + }, + "Request13": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1470,6 +1548,22 @@ ] }, "Request2": { + "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", + "type": "object", + "properties": { + "method": { + "$ref": "#/definitions/DiscoverRequestMethod" + }, + "params": { + "$ref": "#/definitions/DiscoverRequestParams" + } + }, + "required": [ + "method", + "params" + ] + }, + "Request3": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1485,7 +1579,7 @@ "params" ] }, - "Request3": { + "Request4": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1501,7 +1595,7 @@ "params" ] }, - "Request4": { + "Request5": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1517,7 +1611,7 @@ "params" ] }, - "Request5": { + "Request6": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1533,7 +1627,23 @@ "params" ] }, - "Request6": { + "Request7": { + "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", + "type": "object", + "properties": { + "method": { + "$ref": "#/definitions/SubscriptionsListenRequestMethod" + }, + "params": { + "$ref": "#/definitions/SubscriptionsListenRequestParams" + } + }, + "required": [ + "method", + "params" + ] + }, + "Request8": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1549,7 +1659,7 @@ "params" ] }, - "Request7": { + "Request9": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1565,37 +1675,27 @@ "params" ] }, - "Request8": { - "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", + "RequestMetaObject": { + "description": "Metadata reserved by MCP on requests. Extension keys are also allowed.", "type": "object", "properties": { - "method": { - "$ref": "#/definitions/CallToolRequestMethod" + "io.modelcontextprotocol/clientCapabilities": { + "$ref": "#/definitions/ClientCapabilities" }, - "params": { - "$ref": "#/definitions/CallToolRequestParams" + "io.modelcontextprotocol/clientInfo": { + "$ref": "#/definitions/Implementation" + }, + "io.modelcontextprotocol/logLevel": { + "$ref": "#/definitions/LoggingLevel" + }, + "io.modelcontextprotocol/protocolVersion": { + "type": "string" + }, + "progressToken": { + "$ref": "#/definitions/ProgressToken" } }, - "required": [ - "method", - "params" - ] - }, - "Request9": { - "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", - "type": "object", - "properties": { - "method": { - "$ref": "#/definitions/GetTaskMethod" - }, - "params": { - "$ref": "#/definitions/GetTaskParams" - } - }, - "required": [ - "method", - "params" - ] + "additionalProperties": true }, "RequestNoParam": { "type": "object", @@ -1692,38 +1792,20 @@ "method" ] }, - "RequestOptionalParam5": { - "type": "object", - "properties": { - "method": { - "$ref": "#/definitions/ListTasksMethod" - }, - "params": { - "anyOf": [ - { - "$ref": "#/definitions/PaginatedRequestParams" - }, - { - "type": "null" - } - ] - } - }, - "required": [ - "method" - ] - }, "Resource": { "description": "A known resource that the server is capable of reading (spec `Resource`).\n\nAlso used as the inner type of `ContentBlock::ResourceLink` (spec `ResourceLink extends Resource`).", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this resource.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this resource.", @@ -1797,11 +1879,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "mimeType": { "type": [ @@ -1825,11 +1910,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "blob": { "type": "string" @@ -1881,11 +1969,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "name": { "type": [ @@ -2048,19 +2139,6 @@ } ] }, - "SamplingTaskCapability": { - "description": "Sampling task capability. Deprecated by SEP-2577; remains functional and\nwill be removed in a future release.\nSee .", - "type": "object", - "properties": { - "createMessage": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "SetLevelRequestMethod": { "type": "string", "format": "const", @@ -2072,11 +2150,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "level": { "description": "The desired logging level", @@ -2103,11 +2184,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "uri": { "description": "The URI of the resource to subscribe to", @@ -2118,196 +2202,90 @@ "uri" ] }, - "TaskMetadata": { - "description": "Metadata for augmenting a request with task execution (spec `TaskMetadata`).", + "SubscriptionFilter": { + "description": "Notification categories a client opts in to on a `subscriptions/listen` stream.", "type": "object", "properties": { - "ttl": { - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 + "promptsListChanged": { + "type": "boolean" + }, + "resourceSubscriptions": { + "type": "array", + "items": { + "type": "string" + } + }, + "resourcesListChanged": { + "type": "boolean" + }, + "toolsListChanged": { + "type": "boolean" } } }, - "TaskRequestsCapability": { - "description": "Request types that support task-augmented execution.", - "type": "object", - "properties": { - "elicitation": { - "anyOf": [ - { - "$ref": "#/definitions/ElicitationTaskCapability" - }, - { - "type": "null" - } - ] - }, - "sampling": { - "anyOf": [ - { - "$ref": "#/definitions/SamplingTaskCapability" - }, - { - "type": "null" - } - ] - }, - "tools": { - "anyOf": [ - { - "$ref": "#/definitions/ToolsTaskCapability" - }, - { - "type": "null" - } - ] - } - } - }, - "TaskStatus": { - "description": "Canonical task lifecycle status as defined by SEP-1686.", - "oneOf": [ - { - "description": "The receiver accepted the request and is currently working on it.", - "type": "string", - "const": "working" - }, - { - "description": "The receiver requires additional input before work can continue.", - "type": "string", - "const": "input_required" - }, - { - "description": "The underlying operation completed successfully and the result is ready.", - "type": "string", - "const": "completed" - }, - { - "description": "The underlying operation failed and will not continue.", - "type": "string", - "const": "failed" - }, - { - "description": "The task was cancelled and will not continue processing.", - "type": "string", - "const": "cancelled" - } - ] - }, - "TaskStatusNotificationMethod": { + "SubscriptionsListenRequestMethod": { "type": "string", "format": "const", - "const": "notifications/tasks/status" + "const": "subscriptions/listen" }, - "TaskStatusNotificationParam": { - "description": "Parameters for a task status notification (spec `TaskStatusNotificationParams`).\n\nThe task fields are flattened at the top level: `NotificationParams & Task`.", + "SubscriptionsListenRequestParams": { + "description": "Parameters for opening a long-lived notification subscription.", "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "description": "Protocol-level metadata. Required by the 2026-07-28 wire schema.", + "type": "object", + "properties": { + "io.modelcontextprotocol/clientCapabilities": { + "$ref": "#/definitions/ClientCapabilities" + }, + "io.modelcontextprotocol/clientInfo": { + "$ref": "#/definitions/Implementation" + }, + "io.modelcontextprotocol/logLevel": { + "$ref": "#/definitions/LoggingLevel" + }, + "io.modelcontextprotocol/protocolVersion": { + "type": "string" + }, + "progressToken": { + "$ref": "#/definitions/ProgressToken" + } + }, + "additionalProperties": true, + "required": [ + "io.modelcontextprotocol/protocolVersion", + "io.modelcontextprotocol/clientCapabilities" + ] }, - "createdAt": { - "description": "ISO-8601 creation timestamp.", - "type": "string" - }, - "lastUpdatedAt": { - "description": "ISO-8601 timestamp for the most recent status change.", - "type": "string" - }, - "pollInterval": { - "description": "Suggested polling interval (milliseconds).", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 - }, - "status": { - "description": "Current lifecycle status (see [`TaskStatus`]).", + "notifications": { + "description": "Notification categories requested for this stream.", "allOf": [ { - "$ref": "#/definitions/TaskStatus" + "$ref": "#/definitions/SubscriptionFilter" } ] - }, - "statusMessage": { - "description": "Optional human-readable status message for UI surfaces.", - "type": [ - "string", - "null" - ] - }, - "taskId": { - "description": "Unique task identifier generated by the receiver.", - "type": "string" - }, - "ttl": { - "description": "Retention window in milliseconds that the receiver agreed to honor.\n`None` (serialized as `null`) means unlimited retention.", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 } }, "required": [ - "taskId", - "status", - "createdAt", - "lastUpdatedAt" + "_meta", + "notifications" ] }, - "TasksCapability": { - "description": "Task capabilities shared by client and server.", - "type": "object", - "properties": { - "cancel": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "list": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "requests": { - "anyOf": [ - { - "$ref": "#/definitions/TaskRequestsCapability" - }, - { - "type": "null" - } - ] - } - } - }, "TextContent": { "description": "Text content block (spec `TextContent`).", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -2334,11 +2312,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "content": { "type": "array", @@ -2352,13 +2333,7 @@ "null" ] }, - "structuredContent": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, + "structuredContent": true, "toolUseId": { "type": "string" } @@ -2374,11 +2349,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "id": { "type": "string" @@ -2398,18 +2376,6 @@ "input" ] }, - "ToolsTaskCapability": { - "type": "object", - "properties": { - "call": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "UnsubscribeRequestMethod": { "type": "string", "format": "const", @@ -2421,11 +2387,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "uri": { "description": "The URI of the resource to unsubscribe from", @@ -2436,6 +2405,40 @@ "uri" ] }, + "UpdateTaskMethod": { + "type": "string", + "format": "const", + "const": "tasks/update" + }, + "UpdateTaskParams": { + "description": "Parameters for `tasks/update` (SEP-2663): deliver responses to outstanding\nin-task server-to-client requests surfaced via `tasks/get` `inputRequests`.", + "type": "object", + "properties": { + "_meta": { + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] + }, + "inputResponses": { + "description": "Responses to outstanding `inputRequests` previously surfaced by the\nserver. Each key MUST correspond to a currently-outstanding\n`inputRequests` key.", + "type": "object", + "additionalProperties": true + }, + "taskId": { + "description": "Identifier of the task to update.", + "type": "string" + } + }, + "required": [ + "taskId", + "inputResponses" + ] + }, "UrlElicitationCapability": { "description": "Capability for URL mode elicitation.", "type": "object" diff --git a/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema_current.json b/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema_current.json index 3bfdb7c..760d9e1 100644 --- a/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema_current.json +++ b/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema_current.json @@ -54,8 +54,7 @@ "type": [ "string", "null" - ], - "format": "date-time" + ] }, "priority": { "type": [ @@ -87,11 +86,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -124,19 +126,30 @@ "const": "tools/call" }, "CallToolRequestParams": { - "description": "Parameters for calling a tool provided by an MCP server.\n\nContains the tool name and optional arguments needed to execute\nthe tool operation.\n\nThis implements `TaskAugmentedRequestParamsMeta` as tool calls can be\nlong-running and may benefit from task-based execution.", + "description": "Parameters for calling a tool provided by an MCP server.\n\nContains the tool name and optional arguments needed to execute\nthe tool operation.", "type": "object", "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] + }, + "arguments": { + "description": "Arguments to pass to the tool (must match the tool's input schema)", "type": [ "object", "null" ], "additionalProperties": true }, - "arguments": { - "description": "Arguments to pass to the tool (must match the tool's input schema)", + "inputResponses": { + "description": "Client responses to server-initiated input requests from a previous\n[`InputRequiredResult`]. Present only when retrying after an incomplete result.", "type": [ "object", "null" @@ -147,15 +160,11 @@ "description": "The name of the tool to call", "type": "string" }, - "task": { - "description": "Task metadata for async task management (SEP-1319)", - "anyOf": [ - { - "$ref": "#/definitions/TaskMetadata" - }, - { - "type": "null" - } + "requestState": { + "description": "Opaque request state echoed back from a previous [`InputRequiredResult`].\nClients MUST return this value exactly as received.", + "type": [ + "string", + "null" ] } }, @@ -173,11 +182,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "taskId": { "type": "string" @@ -196,11 +208,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "reason": { "type": [ @@ -278,16 +293,6 @@ "type": "null" } ] - }, - "tasks": { - "anyOf": [ - { - "$ref": "#/definitions/TasksCapability" - }, - { - "type": "null" - } - ] } } }, @@ -320,11 +325,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "argument": { "$ref": "#/definitions/ArgumentInfo" @@ -460,11 +468,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "content": { "description": "The actual content of the message (text, image, audio, tool use, or tool result)", @@ -530,17 +541,36 @@ "CustomResult": { "description": "A catch-all response either side can use for custom requests." }, + "DiscoverRequestMethod": { + "type": "string", + "format": "const", + "const": "server/discover" + }, + "DiscoverRequestParams": { + "type": "object", + "properties": { + "_meta": { + "$ref": "#/definitions/RequestMetaObject" + } + }, + "required": [ + "_meta" + ] + }, "ElicitResult": { "description": "The result returned by a client in response to an elicitation request.\n\nContains the user's decision (accept/decline/cancel) and optionally their input data\nif they chose to accept the request.", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this result.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "action": { "description": "The user's decision on how to handle the elicitation request", @@ -606,29 +636,20 @@ } } }, - "ElicitationTaskCapability": { - "type": "object", - "properties": { - "create": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "EmbeddedResource": { "description": "Embedded resource content (spec `EmbeddedResource`).", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -713,11 +734,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "arguments": { "type": [ @@ -726,8 +750,23 @@ ], "additionalProperties": true }, + "inputResponses": { + "description": "Client responses to server-initiated input requests from a previous\n[`InputRequiredResult`].", + "type": [ + "object", + "null" + ], + "additionalProperties": true + }, "name": { "type": "string" + }, + "requestState": { + "description": "Opaque request state echoed back from a previous [`InputRequiredResult`].", + "type": [ + "string", + "null" + ] } }, "required": [ @@ -743,38 +782,17 @@ "type": "object", "properties": { "_meta": { - "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "taskId": { - "type": "string" - } - }, - "required": [ - "taskId" - ] - }, - "GetTaskPayloadMethod": { - "type": "string", - "format": "const", - "const": "tasks/result" - }, - "GetTaskPayloadParams": { - "type": "object", - "properties": { - "_meta": { - "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "taskId": { + "description": "Identifier of the task to query.", "type": "string" } }, @@ -844,11 +862,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -923,11 +944,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "capabilities": { "description": "The capabilities this client supports (sampling, roots, etc.)", @@ -1015,9 +1039,6 @@ { "$ref": "#/definitions/NotificationNoParam2" }, - { - "$ref": "#/definitions/Notification3" - }, { "$ref": "#/definitions/CustomNotification" } @@ -1052,6 +1073,9 @@ { "$ref": "#/definitions/Request4" }, + { + "$ref": "#/definitions/Request5" + }, { "$ref": "#/definitions/RequestOptionalParam" }, @@ -1061,9 +1085,6 @@ { "$ref": "#/definitions/RequestOptionalParam3" }, - { - "$ref": "#/definitions/Request5" - }, { "$ref": "#/definitions/Request6" }, @@ -1073,21 +1094,24 @@ { "$ref": "#/definitions/Request8" }, - { - "$ref": "#/definitions/RequestOptionalParam4" - }, { "$ref": "#/definitions/Request9" }, - { - "$ref": "#/definitions/RequestOptionalParam5" - }, { "$ref": "#/definitions/Request10" }, + { + "$ref": "#/definitions/RequestOptionalParam4" + }, { "$ref": "#/definitions/Request11" }, + { + "$ref": "#/definitions/Request12" + }, + { + "$ref": "#/definitions/Request13" + }, { "$ref": "#/definitions/CustomRequest" } @@ -1140,11 +1164,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "roots": { "type": "array", @@ -1158,11 +1185,6 @@ "roots" ] }, - "ListTasksMethod": { - "type": "string", - "format": "const", - "const": "tasks/list" - }, "ListToolsRequestMethod": { "type": "string", "format": "const", @@ -1183,6 +1205,11 @@ "emergency" ] }, + "MetaObject": { + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", + "type": "object", + "additionalProperties": true + }, "Notification": { "type": "object", "properties": { @@ -1213,20 +1240,15 @@ "params" ] }, - "Notification3": { + "NotificationMetaObject": { + "description": "Metadata reserved by MCP on notifications. Extension keys are also allowed.", "type": "object", "properties": { - "method": { - "$ref": "#/definitions/TaskStatusNotificationMethod" - }, - "params": { - "$ref": "#/definitions/TaskStatusNotificationParam" + "io.modelcontextprotocol/subscriptionId": { + "$ref": "#/definitions/NumberOrString" } }, - "required": [ - "method", - "params" - ] + "additionalProperties": true }, "NotificationNoParam": { "type": "object", @@ -1265,11 +1287,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "cursor": { "type": [ @@ -1293,11 +1318,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "message": { "description": "An optional message describing the current progress.", @@ -1368,12 +1396,30 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] + }, + "inputResponses": { + "description": "Client responses to server-initiated input requests from a previous\n[`InputRequiredResult`].", "type": [ "object", "null" ], "additionalProperties": true }, + "requestState": { + "description": "Opaque request state echoed back from a previous [`InputRequiredResult`].", + "type": [ + "string", + "null" + ] + }, "uri": { "description": "The URI of the resource to read", "type": "string" @@ -1442,10 +1488,10 @@ "type": "object", "properties": { "method": { - "$ref": "#/definitions/GetTaskPayloadMethod" + "$ref": "#/definitions/CallToolRequestMethod" }, "params": { - "$ref": "#/definitions/GetTaskPayloadParams" + "$ref": "#/definitions/CallToolRequestParams" } }, "required": [ @@ -1454,6 +1500,38 @@ ] }, "Request11": { + "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", + "type": "object", + "properties": { + "method": { + "$ref": "#/definitions/GetTaskMethod" + }, + "params": { + "$ref": "#/definitions/GetTaskParams" + } + }, + "required": [ + "method", + "params" + ] + }, + "Request12": { + "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", + "type": "object", + "properties": { + "method": { + "$ref": "#/definitions/UpdateTaskMethod" + }, + "params": { + "$ref": "#/definitions/UpdateTaskParams" + } + }, + "required": [ + "method", + "params" + ] + }, + "Request13": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1470,6 +1548,22 @@ ] }, "Request2": { + "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", + "type": "object", + "properties": { + "method": { + "$ref": "#/definitions/DiscoverRequestMethod" + }, + "params": { + "$ref": "#/definitions/DiscoverRequestParams" + } + }, + "required": [ + "method", + "params" + ] + }, + "Request3": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1485,7 +1579,7 @@ "params" ] }, - "Request3": { + "Request4": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1501,7 +1595,7 @@ "params" ] }, - "Request4": { + "Request5": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1517,7 +1611,7 @@ "params" ] }, - "Request5": { + "Request6": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1533,7 +1627,23 @@ "params" ] }, - "Request6": { + "Request7": { + "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", + "type": "object", + "properties": { + "method": { + "$ref": "#/definitions/SubscriptionsListenRequestMethod" + }, + "params": { + "$ref": "#/definitions/SubscriptionsListenRequestParams" + } + }, + "required": [ + "method", + "params" + ] + }, + "Request8": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1549,7 +1659,7 @@ "params" ] }, - "Request7": { + "Request9": { "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", "type": "object", "properties": { @@ -1565,37 +1675,27 @@ "params" ] }, - "Request8": { - "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", + "RequestMetaObject": { + "description": "Metadata reserved by MCP on requests. Extension keys are also allowed.", "type": "object", "properties": { - "method": { - "$ref": "#/definitions/CallToolRequestMethod" + "io.modelcontextprotocol/clientCapabilities": { + "$ref": "#/definitions/ClientCapabilities" }, - "params": { - "$ref": "#/definitions/CallToolRequestParams" + "io.modelcontextprotocol/clientInfo": { + "$ref": "#/definitions/Implementation" + }, + "io.modelcontextprotocol/logLevel": { + "$ref": "#/definitions/LoggingLevel" + }, + "io.modelcontextprotocol/protocolVersion": { + "type": "string" + }, + "progressToken": { + "$ref": "#/definitions/ProgressToken" } }, - "required": [ - "method", - "params" - ] - }, - "Request9": { - "description": "Represents a JSON-RPC request with method, parameters, and extensions.\n\nThis is the core structure for all MCP requests, containing:\n- `method`: The name of the method being called\n- `params`: The parameters for the method\n- `extensions`: Additional context data (similar to HTTP headers)", - "type": "object", - "properties": { - "method": { - "$ref": "#/definitions/GetTaskMethod" - }, - "params": { - "$ref": "#/definitions/GetTaskParams" - } - }, - "required": [ - "method", - "params" - ] + "additionalProperties": true }, "RequestNoParam": { "type": "object", @@ -1692,38 +1792,20 @@ "method" ] }, - "RequestOptionalParam5": { - "type": "object", - "properties": { - "method": { - "$ref": "#/definitions/ListTasksMethod" - }, - "params": { - "anyOf": [ - { - "$ref": "#/definitions/PaginatedRequestParams" - }, - { - "type": "null" - } - ] - } - }, - "required": [ - "method" - ] - }, "Resource": { "description": "A known resource that the server is capable of reading (spec `Resource`).\n\nAlso used as the inner type of `ContentBlock::ResourceLink` (spec `ResourceLink extends Resource`).", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this resource.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this resource.", @@ -1797,11 +1879,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "mimeType": { "type": [ @@ -1825,11 +1910,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "blob": { "type": "string" @@ -1881,11 +1969,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "name": { "type": [ @@ -2048,19 +2139,6 @@ } ] }, - "SamplingTaskCapability": { - "description": "Sampling task capability. Deprecated by SEP-2577; remains functional and\nwill be removed in a future release.\nSee .", - "type": "object", - "properties": { - "createMessage": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "SetLevelRequestMethod": { "type": "string", "format": "const", @@ -2072,11 +2150,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "level": { "description": "The desired logging level", @@ -2103,11 +2184,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "uri": { "description": "The URI of the resource to subscribe to", @@ -2118,196 +2202,90 @@ "uri" ] }, - "TaskMetadata": { - "description": "Metadata for augmenting a request with task execution (spec `TaskMetadata`).", + "SubscriptionFilter": { + "description": "Notification categories a client opts in to on a `subscriptions/listen` stream.", "type": "object", "properties": { - "ttl": { - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 + "promptsListChanged": { + "type": "boolean" + }, + "resourceSubscriptions": { + "type": "array", + "items": { + "type": "string" + } + }, + "resourcesListChanged": { + "type": "boolean" + }, + "toolsListChanged": { + "type": "boolean" } } }, - "TaskRequestsCapability": { - "description": "Request types that support task-augmented execution.", - "type": "object", - "properties": { - "elicitation": { - "anyOf": [ - { - "$ref": "#/definitions/ElicitationTaskCapability" - }, - { - "type": "null" - } - ] - }, - "sampling": { - "anyOf": [ - { - "$ref": "#/definitions/SamplingTaskCapability" - }, - { - "type": "null" - } - ] - }, - "tools": { - "anyOf": [ - { - "$ref": "#/definitions/ToolsTaskCapability" - }, - { - "type": "null" - } - ] - } - } - }, - "TaskStatus": { - "description": "Canonical task lifecycle status as defined by SEP-1686.", - "oneOf": [ - { - "description": "The receiver accepted the request and is currently working on it.", - "type": "string", - "const": "working" - }, - { - "description": "The receiver requires additional input before work can continue.", - "type": "string", - "const": "input_required" - }, - { - "description": "The underlying operation completed successfully and the result is ready.", - "type": "string", - "const": "completed" - }, - { - "description": "The underlying operation failed and will not continue.", - "type": "string", - "const": "failed" - }, - { - "description": "The task was cancelled and will not continue processing.", - "type": "string", - "const": "cancelled" - } - ] - }, - "TaskStatusNotificationMethod": { + "SubscriptionsListenRequestMethod": { "type": "string", "format": "const", - "const": "notifications/tasks/status" + "const": "subscriptions/listen" }, - "TaskStatusNotificationParam": { - "description": "Parameters for a task status notification (spec `TaskStatusNotificationParams`).\n\nThe task fields are flattened at the top level: `NotificationParams & Task`.", + "SubscriptionsListenRequestParams": { + "description": "Parameters for opening a long-lived notification subscription.", "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "description": "Protocol-level metadata. Required by the 2026-07-28 wire schema.", + "type": "object", + "properties": { + "io.modelcontextprotocol/clientCapabilities": { + "$ref": "#/definitions/ClientCapabilities" + }, + "io.modelcontextprotocol/clientInfo": { + "$ref": "#/definitions/Implementation" + }, + "io.modelcontextprotocol/logLevel": { + "$ref": "#/definitions/LoggingLevel" + }, + "io.modelcontextprotocol/protocolVersion": { + "type": "string" + }, + "progressToken": { + "$ref": "#/definitions/ProgressToken" + } + }, + "additionalProperties": true, + "required": [ + "io.modelcontextprotocol/protocolVersion", + "io.modelcontextprotocol/clientCapabilities" + ] }, - "createdAt": { - "description": "ISO-8601 creation timestamp.", - "type": "string" - }, - "lastUpdatedAt": { - "description": "ISO-8601 timestamp for the most recent status change.", - "type": "string" - }, - "pollInterval": { - "description": "Suggested polling interval (milliseconds).", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 - }, - "status": { - "description": "Current lifecycle status (see [`TaskStatus`]).", + "notifications": { + "description": "Notification categories requested for this stream.", "allOf": [ { - "$ref": "#/definitions/TaskStatus" + "$ref": "#/definitions/SubscriptionFilter" } ] - }, - "statusMessage": { - "description": "Optional human-readable status message for UI surfaces.", - "type": [ - "string", - "null" - ] - }, - "taskId": { - "description": "Unique task identifier generated by the receiver.", - "type": "string" - }, - "ttl": { - "description": "Retention window in milliseconds that the receiver agreed to honor.\n`None` (serialized as `null`) means unlimited retention.", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 } }, "required": [ - "taskId", - "status", - "createdAt", - "lastUpdatedAt" + "_meta", + "notifications" ] }, - "TasksCapability": { - "description": "Task capabilities shared by client and server.", - "type": "object", - "properties": { - "cancel": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "list": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "requests": { - "anyOf": [ - { - "$ref": "#/definitions/TaskRequestsCapability" - }, - { - "type": "null" - } - ] - } - } - }, "TextContent": { "description": "Text content block (spec `TextContent`).", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -2334,11 +2312,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "content": { "type": "array", @@ -2352,13 +2333,7 @@ "null" ] }, - "structuredContent": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, + "structuredContent": true, "toolUseId": { "type": "string" } @@ -2374,11 +2349,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "id": { "type": "string" @@ -2398,18 +2376,6 @@ "input" ] }, - "ToolsTaskCapability": { - "type": "object", - "properties": { - "call": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "UnsubscribeRequestMethod": { "type": "string", "format": "const", @@ -2421,11 +2387,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "uri": { "description": "The URI of the resource to unsubscribe from", @@ -2436,6 +2405,40 @@ "uri" ] }, + "UpdateTaskMethod": { + "type": "string", + "format": "const", + "const": "tasks/update" + }, + "UpdateTaskParams": { + "description": "Parameters for `tasks/update` (SEP-2663): deliver responses to outstanding\nin-task server-to-client requests surfaced via `tasks/get` `inputRequests`.", + "type": "object", + "properties": { + "_meta": { + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] + }, + "inputResponses": { + "description": "Responses to outstanding `inputRequests` previously surfaced by the\nserver. Each key MUST correspond to a currently-outstanding\n`inputRequests` key.", + "type": "object", + "additionalProperties": true + }, + "taskId": { + "description": "Identifier of the task to update.", + "type": "string" + } + }, + "required": [ + "taskId", + "inputResponses" + ] + }, "UrlElicitationCapability": { "description": "Capability for URL mode elicitation.", "type": "object" diff --git a/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json b/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json index fcf821f..e220078 100644 --- a/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json +++ b/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json @@ -54,8 +54,7 @@ "type": [ "string", "null" - ], - "format": "date-time" + ] }, "priority": { "type": [ @@ -77,11 +76,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -151,17 +153,35 @@ "format": "const", "const": "boolean" }, + "CacheScope": { + "description": "Scope describing who may cache cacheable list/read results (SEP-2549).\n\nDefaults to [`CacheScope::Public`] when absent from the wire.", + "oneOf": [ + { + "description": "Any client or intermediary may cache and serve the response to any user.", + "type": "string", + "const": "public" + }, + { + "description": "Only the requesting user's client may cache the response.", + "type": "string", + "const": "private" + } + ] + }, "CallToolResult": { "description": "The result of a tool call operation.\n\nContains the content returned by the tool execution and an optional\nflag indicating whether the operation resulted in an error.", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this result", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "content": { "description": "The content returned by the tool (text, images, etc.)", @@ -178,75 +198,22 @@ "null" ] }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] + }, "structuredContent": { "description": "An optional JSON object that represents the structured result of the tool call" } } }, - "CancelTaskResult": { - "description": "Response to a `tasks/cancel` request.\n\nPer spec, `CancelTaskResult = allOf[Result, Task]` — same shape as `GetTaskResult`.", - "type": "object", - "properties": { - "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "createdAt": { - "description": "ISO-8601 creation timestamp.", - "type": "string" - }, - "lastUpdatedAt": { - "description": "ISO-8601 timestamp for the most recent status change.", - "type": "string" - }, - "pollInterval": { - "description": "Suggested polling interval (milliseconds).", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 - }, - "status": { - "description": "Current lifecycle status (see [`TaskStatus`]).", - "allOf": [ - { - "$ref": "#/definitions/TaskStatus" - } - ] - }, - "statusMessage": { - "description": "Optional human-readable status message for UI surfaces.", - "type": [ - "string", - "null" - ] - }, - "taskId": { - "description": "Unique task identifier generated by the receiver.", - "type": "string" - }, - "ttl": { - "description": "Retention window in milliseconds that the receiver agreed to honor.\n`None` (serialized as `null`) means unlimited retention.", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 - } - }, - "required": [ - "taskId", - "status", - "createdAt", - "lastUpdatedAt" - ] - }, "CancelledNotificationMethod": { "type": "string", "format": "const", @@ -256,11 +223,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "reason": { "type": [ @@ -280,18 +250,93 @@ } } }, - "CompleteResult": { + "ClientCapabilities": { + "title": "Builder", + "description": "```rust\n# use rmcp::model::ClientCapabilities;\nlet cap = ClientCapabilities::builder()\n .enable_experimental()\n .build();\n```", "type": "object", "properties": { - "_meta": { + "elicitation": { + "description": "Capability to handle elicitation requests from servers for interactive user input", + "anyOf": [ + { + "$ref": "#/definitions/ElicitationCapability" + }, + { + "type": "null" + } + ] + }, + "experimental": { "type": [ "object", "null" ], - "additionalProperties": true + "additionalProperties": { + "type": "object", + "additionalProperties": true + } + }, + "extensions": { + "description": "Optional MCP extensions that the client supports (SEP-1724).\nKeys are extension identifiers (e.g., `\"io.modelcontextprotocol/ui\"`),\nvalues are per-extension settings objects. An empty object indicates\nsupport with no settings.", + "type": [ + "object", + "null" + ], + "additionalProperties": { + "type": "object", + "additionalProperties": true + } + }, + "roots": { + "description": "Capability for filesystem roots (deprecated by SEP-2577).", + "anyOf": [ + { + "$ref": "#/definitions/RootsCapabilities" + }, + { + "type": "null" + } + ] + }, + "sampling": { + "description": "Capability for LLM sampling requests (SEP-1577, deprecated by SEP-2577).", + "anyOf": [ + { + "$ref": "#/definitions/SamplingCapability" + }, + { + "type": "null" + } + ] + } + } + }, + "CompleteResult": { + "type": "object", + "properties": { + "_meta": { + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "completion": { "$ref": "#/definitions/CompletionInfo" + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] } }, "required": [ @@ -463,11 +508,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "includeContext": { "description": "How much context to include from MCP servers", @@ -524,17 +572,6 @@ "null" ] }, - "task": { - "description": "Task metadata for async task management (SEP-1319)", - "anyOf": [ - { - "$ref": "#/definitions/TaskMetadata" - }, - { - "type": "null" - } - ] - }, "temperature": { "description": "Temperature for controlling randomness (0.0 to 1.0)", "type": [ @@ -572,22 +609,79 @@ ] }, "CreateTaskResult": { - "description": "Wrapper returned by task-augmented requests (CreateTaskResult in SEP-1686).", + "description": "Result returned in lieu of a standard result to indicate the request will\nbe processed asynchronously (spec `CreateTaskResult`, `resultType: \"task\"`).\n\nThe embedded task is the seed state for the task; the client uses\n`task.task_id` for all subsequent `tasks/get`, `tasks/update`, and\n`tasks/cancel` calls.", "type": "object", "properties": { "_meta": { + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "createdAt": { + "description": "ISO 8601 timestamp when the task was created.", + "type": "string" + }, + "lastUpdatedAt": { + "description": "ISO 8601 timestamp when the task was last updated.", + "type": "string" + }, + "pollIntervalMs": { + "description": "Suggested polling interval in integer milliseconds. Clients SHOULD honor\nthis value to avoid overwhelming the server. This value MAY change over\nthe lifetime of a task.", "type": [ - "object", + "integer", "null" ], - "additionalProperties": true + "format": "uint64", + "minimum": 0 }, - "task": { - "$ref": "#/definitions/Task" + "resultType": { + "description": "Always `\"task\"`.", + "allOf": [ + { + "$ref": "#/definitions/ResultType" + } + ] + }, + "status": { + "description": "Current task status.", + "allOf": [ + { + "$ref": "#/definitions/TaskStatus" + } + ] + }, + "statusMessage": { + "description": "Optional message describing the current task state.\nThis MAY be exposed to the end-user or model.", + "type": [ + "string", + "null" + ] + }, + "taskId": { + "description": "Stable identifier for this task, generated by the server.", + "type": "string" + }, + "ttlMs": { + "description": "Time-to-live duration from creation in integer milliseconds; `None`\n(serialized as `null`) means unlimited. The server may discard the task\nafter the TTL elapses. This value MAY change over the lifetime of a task.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ - "task" + "resultType", + "taskId", + "status", + "createdAt", + "lastUpdatedAt" ] }, "CustomNotification": { @@ -619,6 +713,74 @@ "CustomResult": { "description": "A catch-all response either side can use for custom requests." }, + "DiscoverResult": { + "description": "The server's response to a [`DiscoverRequest`].", + "type": "object", + "properties": { + "_meta": { + "description": "Protocol-level response metadata.", + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Whether the cached result may be shared across authorization contexts.", + "allOf": [ + { + "$ref": "#/definitions/CacheScope" + } + ] + }, + "capabilities": { + "description": "Capabilities provided by this server.", + "allOf": [ + { + "$ref": "#/definitions/ServerCapabilities" + } + ] + }, + "instructions": { + "description": "Optional guidance for using the server.", + "type": [ + "string", + "null" + ] + }, + "resultType": { + "description": "Identifies how the result should be parsed.", + "allOf": [ + { + "$ref": "#/definitions/ResultType" + } + ] + }, + "supportedVersions": { + "description": "Protocol versions implemented by this server.", + "type": "array", + "items": { + "$ref": "#/definitions/ProtocolVersion" + } + }, + "ttlMs": { + "description": "How long clients may consider this response fresh, in milliseconds.", + "type": "integer", + "format": "uint64", + "minimum": 0 + } + }, + "required": [ + "resultType", + "supportedVersions", + "capabilities", + "ttlMs", + "cacheScope" + ] + }, "ElicitRequestParams": { "description": "Parameters for creating an elicitation request to gather user input.\n\nThis structure contains everything needed to request interactive input from a user:\n- A human-readable message explaining what information is needed\n- A type-safe schema defining the expected structure of the response\n\n# Example\n1. Form-based elicitation request\n```rust\nuse rmcp::model::*;\n\nlet params = ElicitRequestParams::FormElicitationParams {\n meta: None,\n message: \"Please provide your email\".to_string(),\n requested_schema: ElicitationSchema::builder()\n .required_email(\"email\")\n .build()\n .unwrap(),\n};\n```\n2. URL-based elicitation request\n```rust\nuse rmcp::model::*;\nlet params = ElicitRequestParams::UrlElicitationParams {\n meta: None,\n message: \"Please provide your feedback at the following URL\".to_string(),\n url: \"https://example.com/feedback\".to_string(),\n elicitation_id: \"unique-id-123\".to_string(),\n};\n```", "anyOf": [ @@ -626,11 +788,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "message": { "type": "string" @@ -653,11 +818,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "elicitationId": { "type": "string" @@ -684,11 +852,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "message": { "type": "string" @@ -710,11 +881,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this result.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "action": { "description": "The user's decision on how to handle the elicitation request", @@ -752,39 +926,47 @@ } ] }, - "ElicitationCompletionNotificationMethod": { - "type": "string", - "format": "const", - "const": "notifications/elicitation/complete" + "ElicitationCapability": { + "description": "Elicitation allows servers to request interactive input from users during tool execution.\nThis capability indicates that a client can handle elicitation requests and present\nappropriate UI to users for collecting the requested information.", + "type": "object", + "properties": { + "form": { + "description": "Whether client supports form-based elicitation.", + "anyOf": [ + { + "$ref": "#/definitions/FormElicitationCapability" + }, + { + "type": "null" + } + ] + }, + "url": { + "description": "Whether client supports URL-based elicitation.", + "anyOf": [ + { + "$ref": "#/definitions/UrlElicitationCapability" + }, + { + "type": "null" + } + ] + } + } }, "ElicitationCreateRequestMethod": { "type": "string", "format": "const", "const": "elicitation/create" }, - "ElicitationResponseNotificationParam": { - "description": "Notification parameters for an url elicitation completion notification.", - "type": "object", - "properties": { - "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "elicitationId": { - "type": "string" - } - }, - "required": [ - "elicitationId" - ] - }, "ElicitationSchema": { "description": "Type-safe elicitation schema for requesting structured user input.\n\nThis enforces the MCP 2025-06-18 specification that elicitation schemas\nmust be objects with primitive-typed properties.\n\n# Example\n\n```rust\nuse rmcp::model::*;\n\nlet schema = ElicitationSchema::builder()\n .required_email(\"email\")\n .required_integer(\"age\", 0, 150)\n .optional_bool(\"newsletter\", false)\n .build();\n```", "type": "object", "properties": { + "$schema": { + "description": "Optional JSON Schema dialect identifier (the `$schema` keyword).\n\nThe 2025-11-25 protocol revision allows a `requestedSchema` to declare its\ndialect. It is preserved verbatim so a declared dialect survives a\ndecode/re-encode round-trip instead of being silently dropped.", + "type": "string" + }, "description": { "description": "Optional description of what this schema represents", "type": [ @@ -830,29 +1012,20 @@ "properties" ] }, - "ElicitationTaskCapability": { - "type": "object", - "properties": { - "create": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "EmbeddedResource": { "description": "Embedded resource content (spec `EmbeddedResource`).", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -927,15 +1100,31 @@ "message" ] }, + "FormElicitationCapability": { + "description": "Capability for handling elicitation requests from servers.\nElicitation allows servers to request interactive input from users during tool execution.\nThis capability indicates that a client can handle elicitation requests and present\nappropriate UI to users for collecting the requested information.\n\nCapability for form mode elicitation.", + "type": "object", + "properties": { + "schemaValidation": { + "description": "Whether the client supports JSON Schema validation for elicitation responses.\nWhen true, the client will validate user input against the requested_schema\nbefore sending the response back to the server.", + "type": [ + "boolean", + "null" + ] + } + } + }, "GetPromptResult": { "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "description": { "type": [ @@ -948,36 +1137,63 @@ "items": { "$ref": "#/definitions/PromptMessage" } + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] } }, "required": [ "messages" ] }, - "GetTaskPayloadResult": { - "description": "Response to a `tasks/result` request.\n\nPer spec, the result structure matches the original request type\n(e.g., `CallToolResult` for `tools/call`). This is represented as\nan open object. The payload is the original request's result\nserialized as a JSON value." - }, "GetTaskResult": { - "description": "Response to a `tasks/get` request.\n\nPer spec, `GetTaskResult = allOf[Result, Task]` — the Task fields are\nflattened at the top level, not nested under a `task` key.", + "description": "Response to a `tasks/get` request (spec `GetTaskResult = Result & DetailedTask`).\n\n`resultType` is `\"complete\"` — this is the standard result shape for\n`tasks/get`, not a task handle.", "type": "object", "properties": { "_meta": { + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "createdAt": { + "description": "ISO 8601 timestamp when the task was created.", + "type": "string" + }, + "error": { "type": [ "object", "null" ], "additionalProperties": true }, - "createdAt": { - "description": "ISO-8601 creation timestamp.", - "type": "string" + "inputRequests": { + "type": [ + "object", + "null" + ], + "additionalProperties": { + "$ref": "#/definitions/InputRequest" + } }, "lastUpdatedAt": { - "description": "ISO-8601 timestamp for the most recent status change.", + "description": "ISO 8601 timestamp when the task was last updated.", "type": "string" }, - "pollInterval": { - "description": "Suggested polling interval (milliseconds).", + "pollIntervalMs": { + "description": "Suggested polling interval in integer milliseconds. Clients SHOULD honor\nthis value to avoid overwhelming the server. This value MAY change over\nthe lifetime of a task.", "type": [ "integer", "null" @@ -985,8 +1201,24 @@ "format": "uint64", "minimum": 0 }, + "result": { + "type": [ + "object", + "null" + ], + "additionalProperties": true + }, + "resultType": { + "description": "Result type discriminator. `tasks/get` responses are standard results:\n`\"complete\"` (SEP-2322). Absent values deserialize as `\"complete\"`.", + "allOf": [ + { + "$ref": "#/definitions/ResultType" + } + ], + "default": "complete" + }, "status": { - "description": "Current lifecycle status (see [`TaskStatus`]).", + "description": "Current task status.", "allOf": [ { "$ref": "#/definitions/TaskStatus" @@ -994,18 +1226,18 @@ ] }, "statusMessage": { - "description": "Optional human-readable status message for UI surfaces.", + "description": "Optional message describing the current task state.\nThis MAY be exposed to the end-user or model.", "type": [ "string", "null" ] }, "taskId": { - "description": "Unique task identifier generated by the receiver.", + "description": "Stable identifier for this task, generated by the server.", "type": "string" }, - "ttl": { - "description": "Retention window in milliseconds that the receiver agreed to honor.\n`None` (serialized as `null`) means unlimited retention.", + "ttlMs": { + "description": "Time-to-live duration from creation in integer milliseconds; `None`\n(serialized as `null`) means unlimited. The server may discard the task\nafter the TTL elapses. This value MAY change over the lifetime of a task.", "type": [ "integer", "null" @@ -1083,11 +1315,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -1161,11 +1396,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "capabilities": { "description": "The capabilities this server provides (tools, resources, prompts, etc.)", @@ -1205,6 +1443,80 @@ "serverInfo" ] }, + "InputRequest": { + "description": "A server-initiated request that can appear inside [`InputRequests`].\n\nPer the MCP spec, only `CreateMessageRequest` (sampling),\n`ElicitRequest` (elicitation), and `ListRootsRequest` (roots)\nare allowed. This is modeled as an untagged enum rather than a\n`ServerRequest` alias to prevent `PingRequest` or `CustomRequest` from\nbeing included.", + "anyOf": [ + { + "description": "A `sampling/createMessage` request.", + "allOf": [ + { + "$ref": "#/definitions/Request" + } + ] + }, + { + "description": "An `elicitation/create` request.", + "allOf": [ + { + "$ref": "#/definitions/Request2" + } + ] + }, + { + "description": "A `roots/list` request.", + "allOf": [ + { + "$ref": "#/definitions/RequestNoParam2" + } + ] + } + ] + }, + "InputRequiredResult": { + "description": "A result indicating that additional input is needed before the request\ncan be completed.\n\nAt least one of [`input_requests`](Self::input_requests) or\n[`request_state`](Self::request_state) MUST be present.\n\nServers MAY send this in response to `tools/call`, `prompts/get`, or\n`resources/read`. Servers MUST NOT send this for any other request.\n\n# Examples\n\n```\nuse rmcp::model::InputRequiredResult;\n\nlet result = InputRequiredResult::from_request_state(\"opaque-server-state\");\nassert!(result.input_requests.is_none());\nassert_eq!(result.request_state.as_deref(), Some(\"opaque-server-state\"));\n```", + "type": "object", + "properties": { + "_meta": { + "description": "Optional protocol-level metadata.", + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "inputRequests": { + "description": "Server-initiated requests that the client must fulfill before retrying.", + "type": [ + "object", + "null" + ], + "additionalProperties": { + "$ref": "#/definitions/InputRequest" + } + }, + "requestState": { + "description": "Opaque request state to be echoed back by the client on retry.\nClients MUST NOT inspect, parse, modify, or make any assumptions\nabout the contents.", + "type": [ + "string", + "null" + ] + }, + "resultType": { + "description": "Always `\"input_required\"` for this result type.", + "allOf": [ + { + "$ref": "#/definitions/ResultType" + } + ] + } + }, + "required": [ + "resultType" + ] + }, "IntegerSchema": { "description": "Schema definition for integer properties.\n\nCompliant with MCP 2025-06-18 specification for elicitation schemas.\nSupports only the fields allowed by the MCP spec.", "type": "object", @@ -1439,11 +1751,25 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Scope describing who may cache this result (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "anyOf": [ + { + "$ref": "#/definitions/CacheScope" + }, + { + "type": "null" + } + ] }, "nextCursor": { "type": [ @@ -1456,6 +1782,26 @@ "items": { "$ref": "#/definitions/Prompt" } + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] + }, + "ttlMs": { + "description": "Time, in milliseconds, that this result may be treated as fresh (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ @@ -1466,11 +1812,25 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Scope describing who may cache this result (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "anyOf": [ + { + "$ref": "#/definitions/CacheScope" + }, + { + "type": "null" + } + ] }, "nextCursor": { "type": [ @@ -1483,6 +1843,26 @@ "items": { "$ref": "#/definitions/ResourceTemplate" } + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] + }, + "ttlMs": { + "description": "Time, in milliseconds, that this result may be treated as fresh (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ @@ -1493,11 +1873,25 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Scope describing who may cache this result (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "anyOf": [ + { + "$ref": "#/definitions/CacheScope" + }, + { + "type": "null" + } + ] }, "nextCursor": { "type": [ @@ -1510,6 +1904,26 @@ "items": { "$ref": "#/definitions/Resource" } + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] + }, + "ttlMs": { + "description": "Time, in milliseconds, that this result may be treated as fresh (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ @@ -1521,15 +1935,29 @@ "format": "const", "const": "roots/list" }, - "ListTasksResult": { + "ListToolsResult": { "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Scope describing who may cache this result (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "anyOf": [ + { + "$ref": "#/definitions/CacheScope" + }, + { + "type": "null" + } + ] }, "nextCursor": { "type": [ @@ -1537,31 +1965,15 @@ "null" ] }, - "tasks": { - "type": "array", - "items": { - "$ref": "#/definitions/Task" - } - } - }, - "required": [ - "tasks" - ] - }, - "ListToolsResult": { - "type": "object", - "properties": { - "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "nextCursor": { - "type": [ - "string", - "null" + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } ] }, "tools": { @@ -1569,6 +1981,15 @@ "items": { "$ref": "#/definitions/Tool" } + }, + "ttlMs": { + "description": "Time, in milliseconds, that this result may be treated as fresh (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ @@ -1600,11 +2021,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "data": { "description": "The actual log data" @@ -1631,6 +2055,11 @@ "data" ] }, + "MetaObject": { + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", + "type": "object", + "additionalProperties": true + }, "ModelHint": { "description": "A hint suggesting a preferred model name or family.\n\nModel hints are advisory suggestions that help clients choose appropriate\nmodels. They can be specific model names or general families like \"claude\" or \"gpt\".", "type": "object", @@ -1761,10 +2190,10 @@ "type": "object", "properties": { "method": { - "$ref": "#/definitions/ElicitationCompletionNotificationMethod" + "$ref": "#/definitions/SubscriptionsAcknowledgedNotificationMethod" }, "params": { - "$ref": "#/definitions/ElicitationResponseNotificationParam" + "$ref": "#/definitions/SubscriptionsAcknowledgedNotificationParams" } }, "required": [ @@ -1779,7 +2208,7 @@ "$ref": "#/definitions/TaskStatusNotificationMethod" }, "params": { - "$ref": "#/definitions/TaskStatusNotificationParam" + "$ref": "#/definitions/TaskStatusNotificationParams" } }, "required": [ @@ -1787,6 +2216,16 @@ "params" ] }, + "NotificationMetaObject": { + "description": "Metadata reserved by MCP on notifications. Extension keys are also allowed.", + "type": "object", + "properties": { + "io.modelcontextprotocol/subscriptionId": { + "$ref": "#/definitions/NumberOrString" + } + }, + "additionalProperties": true + }, "NotificationNoParam": { "type": "object", "properties": { @@ -1954,11 +2393,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "message": { "description": "An optional message describing the current progress.", @@ -2002,11 +2444,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "arguments": { "type": [ @@ -2117,11 +2562,25 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Scope describing who may cache this result (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "anyOf": [ + { + "$ref": "#/definitions/CacheScope" + }, + { + "type": "null" + } + ] }, "contents": { "description": "The actual content of the resource", @@ -2129,6 +2588,26 @@ "items": { "$ref": "#/definitions/ResourceContents" } + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] + }, + "ttlMs": { + "description": "Time, in milliseconds, that this result may be treated as fresh (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ @@ -2167,6 +2646,28 @@ "params" ] }, + "RequestMetaObject": { + "description": "Metadata reserved by MCP on requests. Extension keys are also allowed.", + "type": "object", + "properties": { + "io.modelcontextprotocol/clientCapabilities": { + "$ref": "#/definitions/ClientCapabilities" + }, + "io.modelcontextprotocol/clientInfo": { + "$ref": "#/definitions/Implementation" + }, + "io.modelcontextprotocol/logLevel": { + "$ref": "#/definitions/LoggingLevel" + }, + "io.modelcontextprotocol/protocolVersion": { + "type": "string" + }, + "progressToken": { + "$ref": "#/definitions/ProgressToken" + } + }, + "additionalProperties": true + }, "RequestNoParam": { "type": "object", "properties": { @@ -2195,11 +2696,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this resource.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this resource.", @@ -2273,11 +2777,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "mimeType": { "type": [ @@ -2301,11 +2808,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "blob": { "type": "string" @@ -2338,11 +2848,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this resource template.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this template.", @@ -2410,11 +2923,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "uri": { "description": "The URI of the resource that was updated", @@ -2442,6 +2958,10 @@ } } }, + "ResultType": { + "description": "Indicates the type of a result object, allowing the client to\ndetermine how to parse the response.\n\nThe spec defines this as an open string (`\"complete\" | \"input_required\" | string`),\nso unknown values are preserved rather than rejected. Servers implementing this\nprotocol version MUST include `resultType` in every result. For backward\ncompatibility, clients MUST treat an absent field as `\"complete\"`.\n\nOrdinary results model the field as `Option`: `None` means the\nfield is absent on the wire. Constructors default to `Some(COMPLETE)`, and\nthe server handler strips the `\"complete\"` discriminator before responding\nto peers that negotiated a protocol version older than `2026-07-28`, so\nlegacy sessions keep their historical wire shape (see\n[`ServerResult::strip_result_type_for_legacy_peer`]).", + "type": "string" + }, "Role": { "description": "Represents the role of a participant in a conversation or message exchange.\n\nUsed in sampling and chat contexts to distinguish between different\ntypes of message senders in the conversation flow.", "oneOf": [ @@ -2457,6 +2977,40 @@ } ] }, + "RootsCapabilities": { + "description": "Roots capability. Deprecated by SEP-2577; remains functional and will be\nremoved in a future release.\nSee .", + "type": "object", + "properties": { + "listChanged": { + "type": [ + "boolean", + "null" + ] + } + } + }, + "SamplingCapability": { + "description": "Sampling capability with optional sub-capabilities (SEP-1577).\n\nDeprecated by SEP-2577; remains functional and will be removed in a future\nrelease.\nSee .", + "type": "object", + "properties": { + "context": { + "description": "Support for `includeContext` (soft-deprecated)", + "type": [ + "object", + "null" + ], + "additionalProperties": true + }, + "tools": { + "description": "Support for `tools` and `toolChoice` parameters", + "type": [ + "object", + "null" + ], + "additionalProperties": true + } + } + }, "SamplingContent": { "description": "Single or array content wrapper (SEP-1577).", "anyOf": [ @@ -2476,11 +3030,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "content": { "description": "The actual content of the message (text, image, audio, tool use, or tool result)", @@ -2598,19 +3155,6 @@ } ] }, - "SamplingTaskCapability": { - "description": "Sampling task capability. Deprecated by SEP-2577; remains functional and\nwill be removed in a future release.\nSee .", - "type": "object", - "properties": { - "createMessage": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "ServerCapabilities": { "title": "Builder", "description": "```rust\n# use rmcp::model::ServerCapabilities;\nlet cap = ServerCapabilities::builder()\n .enable_experimental()\n .enable_prompts()\n .enable_resources()\n .enable_tools()\n .enable_tool_list_changed()\n .build();\n```", @@ -2672,16 +3216,6 @@ } ] }, - "tasks": { - "anyOf": [ - { - "$ref": "#/definitions/TasksCapability" - }, - { - "type": "null" - } - ] - }, "tools": { "anyOf": [ { @@ -2696,6 +3230,9 @@ }, "ServerResult": { "anyOf": [ + { + "$ref": "#/definitions/DiscoverResult" + }, { "$ref": "#/definitions/InitializeResult" }, @@ -2717,6 +3254,9 @@ { "$ref": "#/definitions/ReadResourceResult" }, + { + "$ref": "#/definitions/SubscriptionsListenResult" + }, { "$ref": "#/definitions/ListToolsResult" }, @@ -2726,20 +3266,17 @@ { "$ref": "#/definitions/CreateTaskResult" }, - { - "$ref": "#/definitions/ListTasksResult" - }, { "$ref": "#/definitions/GetTaskResult" }, - { - "$ref": "#/definitions/CancelTaskResult" - }, { "$ref": "#/definitions/CallToolResult" }, { - "$ref": "#/definitions/GetTaskPayloadResult" + "$ref": "#/definitions/InputRequiredResult" + }, + { + "$ref": "#/definitions/TaskAckResult" }, { "$ref": "#/definitions/EmptyObject" @@ -2857,138 +3394,135 @@ "format": "const", "const": "string" }, - "Task": { - "description": "Primary Task object that surfaces metadata during the task lifecycle.\n\nPer spec, `lastUpdatedAt` and `ttl` are required fields.\n`ttl` is nullable (`null` means unlimited retention).", + "SubscriptionFilter": { + "description": "Notification categories a client opts in to on a `subscriptions/listen` stream.", "type": "object", "properties": { - "createdAt": { - "description": "ISO-8601 creation timestamp.", - "type": "string" + "promptsListChanged": { + "type": "boolean" }, - "lastUpdatedAt": { - "description": "ISO-8601 timestamp for the most recent status change.", - "type": "string" + "resourceSubscriptions": { + "type": "array", + "items": { + "type": "string" + } }, - "pollInterval": { - "description": "Suggested polling interval (milliseconds).", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 + "resourcesListChanged": { + "type": "boolean" }, - "status": { - "description": "Current lifecycle status (see [`TaskStatus`]).", - "allOf": [ - { - "$ref": "#/definitions/TaskStatus" - } - ] + "toolsListChanged": { + "type": "boolean" + } + } + }, + "SubscriptionsAcknowledgedNotificationMethod": { + "type": "string", + "format": "const", + "const": "notifications/subscriptions/acknowledged" + }, + "SubscriptionsAcknowledgedNotificationParams": { + "description": "Parameters reporting the accepted subset of a subscription filter.", + "type": "object", + "properties": { + "_meta": { + "$ref": "#/definitions/NotificationMetaObject" }, - "statusMessage": { - "description": "Optional human-readable status message for UI surfaces.", - "type": [ - "string", - "null" - ] - }, - "taskId": { - "description": "Unique task identifier generated by the receiver.", - "type": "string" - }, - "ttl": { - "description": "Retention window in milliseconds that the receiver agreed to honor.\n`None` (serialized as `null`) means unlimited retention.", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 + "notifications": { + "$ref": "#/definitions/SubscriptionFilter" } }, "required": [ - "taskId", - "status", - "createdAt", - "lastUpdatedAt" + "notifications" ] }, - "TaskMetadata": { - "description": "Metadata for augmenting a request with task execution (spec `TaskMetadata`).", + "SubscriptionsListenResult": { + "description": "Final response indicating that a subscription ended gracefully.", "type": "object", "properties": { - "ttl": { - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 + "_meta": { + "$ref": "#/definitions/SubscriptionsListenResultMeta" + }, + "resultType": { + "$ref": "#/definitions/ResultType" } - } + }, + "required": [ + "resultType", + "_meta" + ] }, - "TaskRequestsCapability": { - "description": "Request types that support task-augmented execution.", + "SubscriptionsListenResultMeta": { "type": "object", "properties": { - "elicitation": { + "io.modelcontextprotocol/serverInfo": { + "description": "Identifies the server software producing the response. Servers SHOULD include this field on every response unless specifically configured not to do so.", + "allOf": [ + { + "$ref": "#/definitions/Implementation" + } + ] + }, + "io.modelcontextprotocol/subscriptionId": { + "$ref": "#/definitions/NumberOrString" + } + }, + "additionalProperties": true, + "required": [ + "io.modelcontextprotocol/subscriptionId" + ] + }, + "TaskAckResult": { + "description": "Empty acknowledgement for `tasks/update` and `tasks/cancel` (SEP-2663).\n\nThe spec requires these acks to be empty results carrying the SEP-2322\n`resultType: \"complete\"` discriminator; task state changes are observed\nvia the next `tasks/get`.", + "type": "object", + "properties": { + "_meta": { "anyOf": [ { - "$ref": "#/definitions/ElicitationTaskCapability" + "$ref": "#/definitions/MetaObject" }, { "type": "null" } ] }, - "sampling": { - "anyOf": [ + "resultType": { + "description": "Always `\"complete\"`.", + "allOf": [ { - "$ref": "#/definitions/SamplingTaskCapability" - }, - { - "type": "null" - } - ] - }, - "tools": { - "anyOf": [ - { - "$ref": "#/definitions/ToolsTaskCapability" - }, - { - "type": "null" + "$ref": "#/definitions/ResultType" } ] } - } + }, + "required": [ + "resultType" + ] }, "TaskStatus": { - "description": "Canonical task lifecycle status as defined by SEP-1686.", + "description": "Canonical task lifecycle status (SEP-2663).", "oneOf": [ { - "description": "The receiver accepted the request and is currently working on it.", + "description": "The request is currently being processed.", "type": "string", "const": "working" }, { - "description": "The receiver requires additional input before work can continue.", + "description": "The server needs input from the client before the task can proceed.", "type": "string", "const": "input_required" }, { - "description": "The underlying operation completed successfully and the result is ready.", + "description": "The request completed successfully and the result is available.\nThis includes tool calls that returned results with `isError: true`.", "type": "string", "const": "completed" }, { - "description": "The underlying operation failed and will not continue.", + "description": "The request failed due to a JSON-RPC error during execution.", "type": "string", "const": "failed" }, { - "description": "The task was cancelled and will not continue processing.", + "description": "The request was cancelled before completion.", "type": "string", "const": "cancelled" } @@ -2997,29 +3531,48 @@ "TaskStatusNotificationMethod": { "type": "string", "format": "const", - "const": "notifications/tasks/status" + "const": "notifications/tasks" }, - "TaskStatusNotificationParam": { - "description": "Parameters for a task status notification (spec `TaskStatusNotificationParams`).\n\nThe task fields are flattened at the top level: `NotificationParams & Task`.", + "TaskStatusNotificationParams": { + "description": "Parameters for a task status notification (spec `TaskStatusNotificationParams`).\n\nCarries a complete [`DetailedTask`] for the current status, identical to\nwhat `tasks/get` would have returned at that moment. The task fields are\nflattened at the top level: `NotificationParams & Task`.", "type": "object", "properties": { "_meta": { + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] + }, + "createdAt": { + "description": "ISO 8601 timestamp when the task was created.", + "type": "string" + }, + "error": { "type": [ "object", "null" ], "additionalProperties": true }, - "createdAt": { - "description": "ISO-8601 creation timestamp.", - "type": "string" + "inputRequests": { + "type": [ + "object", + "null" + ], + "additionalProperties": { + "$ref": "#/definitions/InputRequest" + } }, "lastUpdatedAt": { - "description": "ISO-8601 timestamp for the most recent status change.", + "description": "ISO 8601 timestamp when the task was last updated.", "type": "string" }, - "pollInterval": { - "description": "Suggested polling interval (milliseconds).", + "pollIntervalMs": { + "description": "Suggested polling interval in integer milliseconds. Clients SHOULD honor\nthis value to avoid overwhelming the server. This value MAY change over\nthe lifetime of a task.", "type": [ "integer", "null" @@ -3027,8 +3580,15 @@ "format": "uint64", "minimum": 0 }, + "result": { + "type": [ + "object", + "null" + ], + "additionalProperties": true + }, "status": { - "description": "Current lifecycle status (see [`TaskStatus`]).", + "description": "Current task status.", "allOf": [ { "$ref": "#/definitions/TaskStatus" @@ -3036,18 +3596,18 @@ ] }, "statusMessage": { - "description": "Optional human-readable status message for UI surfaces.", + "description": "Optional message describing the current task state.\nThis MAY be exposed to the end-user or model.", "type": [ "string", "null" ] }, "taskId": { - "description": "Unique task identifier generated by the receiver.", + "description": "Stable identifier for this task, generated by the server.", "type": "string" }, - "ttl": { - "description": "Retention window in milliseconds that the receiver agreed to honor.\n`None` (serialized as `null`) means unlimited retention.", + "ttlMs": { + "description": "Time-to-live duration from creation in integer milliseconds; `None`\n(serialized as `null`) means unlimited. The server may discard the task\nafter the TTL elapses. This value MAY change over the lifetime of a task.", "type": [ "integer", "null" @@ -3063,67 +3623,20 @@ "lastUpdatedAt" ] }, - "TaskSupport": { - "description": "Per-tool task support mode as defined in the MCP specification.\n\nThis enum indicates whether a tool supports task-based invocation,\nallowing clients to know how to properly call the tool.\n\nSee [Tool-Level Negotiation](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks#tool-level-negotiation).", - "oneOf": [ - { - "description": "Clients MUST NOT invoke this tool as a task (default behavior).", - "type": "string", - "const": "forbidden" - }, - { - "description": "Clients MAY invoke this tool as either a task or a normal call.", - "type": "string", - "const": "optional" - }, - { - "description": "Clients MUST invoke this tool as a task.", - "type": "string", - "const": "required" - } - ] - }, - "TasksCapability": { - "description": "Task capabilities shared by client and server.", - "type": "object", - "properties": { - "cancel": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "list": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "requests": { - "anyOf": [ - { - "$ref": "#/definitions/TaskRequestsCapability" - }, - { - "type": "null" - } - ] - } - } - }, "TextContent": { "description": "Text content block (spec `TextContent`).", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -3256,11 +3769,14 @@ "properties": { "_meta": { "description": "Optional additional metadata for this tool", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional additional tool information.", @@ -3280,17 +3796,6 @@ "null" ] }, - "execution": { - "description": "Execution-related configuration including task support mode.", - "anyOf": [ - { - "$ref": "#/definitions/ToolExecution" - }, - { - "type": "null" - } - ] - }, "icons": { "description": "Optional list of icons for the tool", "type": [ @@ -3409,23 +3914,6 @@ } ] }, - "ToolExecution": { - "description": "Execution-related configuration for a tool.\n\nThis struct contains settings that control how a tool should be executed,\nincluding task support configuration.", - "type": "object", - "properties": { - "taskSupport": { - "description": "Indicates whether this tool supports task-based invocation.\n\nWhen not present or set to `Forbidden`, clients MUST NOT invoke this tool as a task.\nWhen set to `Optional`, clients MAY invoke this tool as a task or normal call.\nWhen set to `Required`, clients MUST invoke this tool as a task.", - "anyOf": [ - { - "$ref": "#/definitions/TaskSupport" - }, - { - "type": "null" - } - ] - } - } - }, "ToolListChangedNotificationMethod": { "type": "string", "format": "const", @@ -3436,11 +3924,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "content": { "type": "array", @@ -3454,13 +3945,7 @@ "null" ] }, - "structuredContent": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, + "structuredContent": true, "toolUseId": { "type": "string" } @@ -3476,11 +3961,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "id": { "type": "string" @@ -3511,18 +3999,6 @@ } } }, - "ToolsTaskCapability": { - "type": "object", - "properties": { - "call": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "UntitledItems": { "description": "Items for untitled multi-select options", "type": "object", @@ -3632,6 +4108,10 @@ "type", "enum" ] + }, + "UrlElicitationCapability": { + "description": "Capability for URL mode elicitation.", + "type": "object" } } } \ No newline at end of file diff --git a/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema_current.json b/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema_current.json index fcf821f..ca0ef56 100644 --- a/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema_current.json +++ b/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema_current.json @@ -54,8 +54,7 @@ "type": [ "string", "null" - ], - "format": "date-time" + ] }, "priority": { "type": [ @@ -77,11 +76,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -151,17 +153,35 @@ "format": "const", "const": "boolean" }, + "CacheScope": { + "description": "Scope describing who may cache cacheable list/read results (SEP-2549).\n\nDefaults to [`CacheScope::Public`] when absent from the wire.", + "oneOf": [ + { + "description": "Any client or intermediary may cache and serve the response to any user.", + "type": "string", + "const": "public" + }, + { + "description": "Only the requesting user's client may cache the response.", + "type": "string", + "const": "private" + } + ] + }, "CallToolResult": { "description": "The result of a tool call operation.\n\nContains the content returned by the tool execution and an optional\nflag indicating whether the operation resulted in an error.", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this result", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "content": { "description": "The content returned by the tool (text, images, etc.)", @@ -178,75 +198,22 @@ "null" ] }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] + }, "structuredContent": { "description": "An optional JSON object that represents the structured result of the tool call" } } }, - "CancelTaskResult": { - "description": "Response to a `tasks/cancel` request.\n\nPer spec, `CancelTaskResult = allOf[Result, Task]` — same shape as `GetTaskResult`.", - "type": "object", - "properties": { - "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "createdAt": { - "description": "ISO-8601 creation timestamp.", - "type": "string" - }, - "lastUpdatedAt": { - "description": "ISO-8601 timestamp for the most recent status change.", - "type": "string" - }, - "pollInterval": { - "description": "Suggested polling interval (milliseconds).", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 - }, - "status": { - "description": "Current lifecycle status (see [`TaskStatus`]).", - "allOf": [ - { - "$ref": "#/definitions/TaskStatus" - } - ] - }, - "statusMessage": { - "description": "Optional human-readable status message for UI surfaces.", - "type": [ - "string", - "null" - ] - }, - "taskId": { - "description": "Unique task identifier generated by the receiver.", - "type": "string" - }, - "ttl": { - "description": "Retention window in milliseconds that the receiver agreed to honor.\n`None` (serialized as `null`) means unlimited retention.", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 - } - }, - "required": [ - "taskId", - "status", - "createdAt", - "lastUpdatedAt" - ] - }, "CancelledNotificationMethod": { "type": "string", "format": "const", @@ -256,11 +223,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "reason": { "type": [ @@ -280,18 +250,93 @@ } } }, - "CompleteResult": { + "ClientCapabilities": { + "title": "Builder", + "description": "```rust\n# use rmcp::model::ClientCapabilities;\nlet cap = ClientCapabilities::builder()\n .enable_experimental()\n .build();\n```", "type": "object", "properties": { - "_meta": { + "elicitation": { + "description": "Capability to handle elicitation requests from servers for interactive user input", + "anyOf": [ + { + "$ref": "#/definitions/ElicitationCapability" + }, + { + "type": "null" + } + ] + }, + "experimental": { "type": [ "object", "null" ], - "additionalProperties": true + "additionalProperties": { + "type": "object", + "additionalProperties": true + } + }, + "extensions": { + "description": "Optional MCP extensions that the client supports (SEP-1724).\nKeys are extension identifiers (e.g., `\"io.modelcontextprotocol/ui\"`),\nvalues are per-extension settings objects. An empty object indicates\nsupport with no settings.", + "type": [ + "object", + "null" + ], + "additionalProperties": { + "type": "object", + "additionalProperties": true + } + }, + "roots": { + "description": "Capability for filesystem roots (deprecated by SEP-2577).", + "anyOf": [ + { + "$ref": "#/definitions/RootsCapabilities" + }, + { + "type": "null" + } + ] + }, + "sampling": { + "description": "Capability for LLM sampling requests (SEP-1577, deprecated by SEP-2577).", + "anyOf": [ + { + "$ref": "#/definitions/SamplingCapability" + }, + { + "type": "null" + } + ] + } + } + }, + "CompleteResult": { + "type": "object", + "properties": { + "_meta": { + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "completion": { "$ref": "#/definitions/CompletionInfo" + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] } }, "required": [ @@ -463,11 +508,14 @@ "properties": { "_meta": { "description": "Protocol-level metadata for this request (SEP-1319)", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "includeContext": { "description": "How much context to include from MCP servers", @@ -524,17 +572,6 @@ "null" ] }, - "task": { - "description": "Task metadata for async task management (SEP-1319)", - "anyOf": [ - { - "$ref": "#/definitions/TaskMetadata" - }, - { - "type": "null" - } - ] - }, "temperature": { "description": "Temperature for controlling randomness (0.0 to 1.0)", "type": [ @@ -572,22 +609,79 @@ ] }, "CreateTaskResult": { - "description": "Wrapper returned by task-augmented requests (CreateTaskResult in SEP-1686).", + "description": "Result returned in lieu of a standard result to indicate the request will\nbe processed asynchronously (spec `CreateTaskResult`, `resultType: \"task\"`).\n\nThe embedded task is the seed state for the task; the client uses\n`task.task_id` for all subsequent `tasks/get`, `tasks/update`, and\n`tasks/cancel` calls.", "type": "object", "properties": { "_meta": { + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "createdAt": { + "description": "ISO 8601 timestamp when the task was created.", + "type": "string" + }, + "lastUpdatedAt": { + "description": "ISO 8601 timestamp when the task was last updated.", + "type": "string" + }, + "pollIntervalMs": { + "description": "Suggested polling interval in integer milliseconds. Clients SHOULD honor\nthis value to avoid overwhelming the server. This value MAY change over\nthe lifetime of a task.", "type": [ - "object", + "integer", "null" ], - "additionalProperties": true + "format": "uint64", + "minimum": 0 }, - "task": { - "$ref": "#/definitions/Task" + "resultType": { + "description": "Always `\"task\"`.", + "allOf": [ + { + "$ref": "#/definitions/ResultType" + } + ] + }, + "status": { + "description": "Current task status.", + "allOf": [ + { + "$ref": "#/definitions/TaskStatus" + } + ] + }, + "statusMessage": { + "description": "Optional message describing the current task state.\nThis MAY be exposed to the end-user or model.", + "type": [ + "string", + "null" + ] + }, + "taskId": { + "description": "Stable identifier for this task, generated by the server.", + "type": "string" + }, + "ttlMs": { + "description": "Time-to-live duration from creation in integer milliseconds; `None`\n(serialized as `null`) means unlimited. The server may discard the task\nafter the TTL elapses. This value MAY change over the lifetime of a task.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ - "task" + "resultType", + "taskId", + "status", + "createdAt", + "lastUpdatedAt" ] }, "CustomNotification": { @@ -619,6 +713,74 @@ "CustomResult": { "description": "A catch-all response either side can use for custom requests." }, + "DiscoverResult": { + "description": "The server's response to a [`DiscoverRequest`].", + "type": "object", + "properties": { + "_meta": { + "description": "Protocol-level response metadata.", + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Whether the cached result may be shared across authorization contexts.", + "allOf": [ + { + "$ref": "#/definitions/CacheScope" + } + ] + }, + "capabilities": { + "description": "Capabilities provided by this server.", + "allOf": [ + { + "$ref": "#/definitions/ServerCapabilities" + } + ] + }, + "instructions": { + "description": "Optional guidance for using the server.", + "type": [ + "string", + "null" + ] + }, + "resultType": { + "description": "Identifies how the result should be parsed.", + "allOf": [ + { + "$ref": "#/definitions/ResultType" + } + ] + }, + "supportedVersions": { + "description": "Protocol versions implemented by this server.", + "type": "array", + "items": { + "$ref": "#/definitions/ProtocolVersion" + } + }, + "ttlMs": { + "description": "How long clients may consider this response fresh, in milliseconds.", + "type": "integer", + "format": "uint64", + "minimum": 0 + } + }, + "required": [ + "resultType", + "supportedVersions", + "capabilities", + "ttlMs", + "cacheScope" + ] + }, "ElicitRequestParams": { "description": "Parameters for creating an elicitation request to gather user input.\n\nThis structure contains everything needed to request interactive input from a user:\n- A human-readable message explaining what information is needed\n- A type-safe schema defining the expected structure of the response\n\n# Example\n1. Form-based elicitation request\n```rust\nuse rmcp::model::*;\n\nlet params = ElicitRequestParams::FormElicitationParams {\n meta: None,\n message: \"Please provide your email\".to_string(),\n requested_schema: ElicitationSchema::builder()\n .required_email(\"email\")\n .build()\n .unwrap(),\n};\n```\n2. URL-based elicitation request\n```rust\nuse rmcp::model::*;\nlet params = ElicitRequestParams::UrlElicitationParams {\n meta: None,\n message: \"Please provide your feedback at the following URL\".to_string(),\n url: \"https://example.com/feedback\".to_string(),\n elicitation_id: \"unique-id-123\".to_string(),\n};\n```", "anyOf": [ @@ -626,11 +788,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "message": { "type": "string" @@ -653,11 +818,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "elicitationId": { "type": "string" @@ -684,11 +852,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/RequestMetaObject" + }, + { + "type": "null" + } + ] }, "message": { "type": "string" @@ -710,11 +881,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this result.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "action": { "description": "The user's decision on how to handle the elicitation request", @@ -752,35 +926,39 @@ } ] }, - "ElicitationCompletionNotificationMethod": { - "type": "string", - "format": "const", - "const": "notifications/elicitation/complete" + "ElicitationCapability": { + "description": "Elicitation allows servers to request interactive input from users during tool execution.\nThis capability indicates that a client can handle elicitation requests and present\nappropriate UI to users for collecting the requested information.", + "type": "object", + "properties": { + "form": { + "description": "Whether client supports form-based elicitation.", + "anyOf": [ + { + "$ref": "#/definitions/FormElicitationCapability" + }, + { + "type": "null" + } + ] + }, + "url": { + "description": "Whether client supports URL-based elicitation.", + "anyOf": [ + { + "$ref": "#/definitions/UrlElicitationCapability" + }, + { + "type": "null" + } + ] + } + } }, "ElicitationCreateRequestMethod": { "type": "string", "format": "const", "const": "elicitation/create" }, - "ElicitationResponseNotificationParam": { - "description": "Notification parameters for an url elicitation completion notification.", - "type": "object", - "properties": { - "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "elicitationId": { - "type": "string" - } - }, - "required": [ - "elicitationId" - ] - }, "ElicitationSchema": { "description": "Type-safe elicitation schema for requesting structured user input.\n\nThis enforces the MCP 2025-06-18 specification that elicitation schemas\nmust be objects with primitive-typed properties.\n\n# Example\n\n```rust\nuse rmcp::model::*;\n\nlet schema = ElicitationSchema::builder()\n .required_email(\"email\")\n .required_integer(\"age\", 0, 150)\n .optional_bool(\"newsletter\", false)\n .build();\n```", "type": "object", @@ -830,29 +1008,20 @@ "properties" ] }, - "ElicitationTaskCapability": { - "type": "object", - "properties": { - "create": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "EmbeddedResource": { "description": "Embedded resource content (spec `EmbeddedResource`).", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -927,15 +1096,31 @@ "message" ] }, + "FormElicitationCapability": { + "description": "Capability for handling elicitation requests from servers.\nElicitation allows servers to request interactive input from users during tool execution.\nThis capability indicates that a client can handle elicitation requests and present\nappropriate UI to users for collecting the requested information.\n\nCapability for form mode elicitation.", + "type": "object", + "properties": { + "schemaValidation": { + "description": "Whether the client supports JSON Schema validation for elicitation responses.\nWhen true, the client will validate user input against the requested_schema\nbefore sending the response back to the server.", + "type": [ + "boolean", + "null" + ] + } + } + }, "GetPromptResult": { "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "description": { "type": [ @@ -948,36 +1133,63 @@ "items": { "$ref": "#/definitions/PromptMessage" } + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] } }, "required": [ "messages" ] }, - "GetTaskPayloadResult": { - "description": "Response to a `tasks/result` request.\n\nPer spec, the result structure matches the original request type\n(e.g., `CallToolResult` for `tools/call`). This is represented as\nan open object. The payload is the original request's result\nserialized as a JSON value." - }, "GetTaskResult": { - "description": "Response to a `tasks/get` request.\n\nPer spec, `GetTaskResult = allOf[Result, Task]` — the Task fields are\nflattened at the top level, not nested under a `task` key.", + "description": "Response to a `tasks/get` request (spec `GetTaskResult = Result & DetailedTask`).\n\n`resultType` is `\"complete\"` — this is the standard result shape for\n`tasks/get`, not a task handle.", "type": "object", "properties": { "_meta": { + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "createdAt": { + "description": "ISO 8601 timestamp when the task was created.", + "type": "string" + }, + "error": { "type": [ "object", "null" ], "additionalProperties": true }, - "createdAt": { - "description": "ISO-8601 creation timestamp.", - "type": "string" + "inputRequests": { + "type": [ + "object", + "null" + ], + "additionalProperties": { + "$ref": "#/definitions/InputRequest" + } }, "lastUpdatedAt": { - "description": "ISO-8601 timestamp for the most recent status change.", + "description": "ISO 8601 timestamp when the task was last updated.", "type": "string" }, - "pollInterval": { - "description": "Suggested polling interval (milliseconds).", + "pollIntervalMs": { + "description": "Suggested polling interval in integer milliseconds. Clients SHOULD honor\nthis value to avoid overwhelming the server. This value MAY change over\nthe lifetime of a task.", "type": [ "integer", "null" @@ -985,8 +1197,24 @@ "format": "uint64", "minimum": 0 }, + "result": { + "type": [ + "object", + "null" + ], + "additionalProperties": true + }, + "resultType": { + "description": "Result type discriminator. `tasks/get` responses are standard results:\n`\"complete\"` (SEP-2322). Absent values deserialize as `\"complete\"`.", + "allOf": [ + { + "$ref": "#/definitions/ResultType" + } + ], + "default": "complete" + }, "status": { - "description": "Current lifecycle status (see [`TaskStatus`]).", + "description": "Current task status.", "allOf": [ { "$ref": "#/definitions/TaskStatus" @@ -994,18 +1222,18 @@ ] }, "statusMessage": { - "description": "Optional human-readable status message for UI surfaces.", + "description": "Optional message describing the current task state.\nThis MAY be exposed to the end-user or model.", "type": [ "string", "null" ] }, "taskId": { - "description": "Unique task identifier generated by the receiver.", + "description": "Stable identifier for this task, generated by the server.", "type": "string" }, - "ttl": { - "description": "Retention window in milliseconds that the receiver agreed to honor.\n`None` (serialized as `null`) means unlimited retention.", + "ttlMs": { + "description": "Time-to-live duration from creation in integer milliseconds; `None`\n(serialized as `null`) means unlimited. The server may discard the task\nafter the TTL elapses. This value MAY change over the lifetime of a task.", "type": [ "integer", "null" @@ -1083,11 +1311,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -1161,11 +1392,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "capabilities": { "description": "The capabilities this server provides (tools, resources, prompts, etc.)", @@ -1205,6 +1439,80 @@ "serverInfo" ] }, + "InputRequest": { + "description": "A server-initiated request that can appear inside [`InputRequests`].\n\nPer the MCP spec, only `CreateMessageRequest` (sampling),\n`ElicitRequest` (elicitation), and `ListRootsRequest` (roots)\nare allowed. This is modeled as an untagged enum rather than a\n`ServerRequest` alias to prevent `PingRequest` or `CustomRequest` from\nbeing included.", + "anyOf": [ + { + "description": "A `sampling/createMessage` request.", + "allOf": [ + { + "$ref": "#/definitions/Request" + } + ] + }, + { + "description": "An `elicitation/create` request.", + "allOf": [ + { + "$ref": "#/definitions/Request2" + } + ] + }, + { + "description": "A `roots/list` request.", + "allOf": [ + { + "$ref": "#/definitions/RequestNoParam2" + } + ] + } + ] + }, + "InputRequiredResult": { + "description": "A result indicating that additional input is needed before the request\ncan be completed.\n\nAt least one of [`input_requests`](Self::input_requests) or\n[`request_state`](Self::request_state) MUST be present.\n\nServers MAY send this in response to `tools/call`, `prompts/get`, or\n`resources/read`. Servers MUST NOT send this for any other request.\n\n# Examples\n\n```\nuse rmcp::model::InputRequiredResult;\n\nlet result = InputRequiredResult::from_request_state(\"opaque-server-state\");\nassert!(result.input_requests.is_none());\nassert_eq!(result.request_state.as_deref(), Some(\"opaque-server-state\"));\n```", + "type": "object", + "properties": { + "_meta": { + "description": "Optional protocol-level metadata.", + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "inputRequests": { + "description": "Server-initiated requests that the client must fulfill before retrying.", + "type": [ + "object", + "null" + ], + "additionalProperties": { + "$ref": "#/definitions/InputRequest" + } + }, + "requestState": { + "description": "Opaque request state to be echoed back by the client on retry.\nClients MUST NOT inspect, parse, modify, or make any assumptions\nabout the contents.", + "type": [ + "string", + "null" + ] + }, + "resultType": { + "description": "Always `\"input_required\"` for this result type.", + "allOf": [ + { + "$ref": "#/definitions/ResultType" + } + ] + } + }, + "required": [ + "resultType" + ] + }, "IntegerSchema": { "description": "Schema definition for integer properties.\n\nCompliant with MCP 2025-06-18 specification for elicitation schemas.\nSupports only the fields allowed by the MCP spec.", "type": "object", @@ -1439,11 +1747,25 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Scope describing who may cache this result (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "anyOf": [ + { + "$ref": "#/definitions/CacheScope" + }, + { + "type": "null" + } + ] }, "nextCursor": { "type": [ @@ -1456,6 +1778,26 @@ "items": { "$ref": "#/definitions/Prompt" } + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] + }, + "ttlMs": { + "description": "Time, in milliseconds, that this result may be treated as fresh (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ @@ -1466,11 +1808,25 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Scope describing who may cache this result (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "anyOf": [ + { + "$ref": "#/definitions/CacheScope" + }, + { + "type": "null" + } + ] }, "nextCursor": { "type": [ @@ -1483,6 +1839,26 @@ "items": { "$ref": "#/definitions/ResourceTemplate" } + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] + }, + "ttlMs": { + "description": "Time, in milliseconds, that this result may be treated as fresh (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ @@ -1493,11 +1869,25 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Scope describing who may cache this result (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "anyOf": [ + { + "$ref": "#/definitions/CacheScope" + }, + { + "type": "null" + } + ] }, "nextCursor": { "type": [ @@ -1510,6 +1900,26 @@ "items": { "$ref": "#/definitions/Resource" } + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] + }, + "ttlMs": { + "description": "Time, in milliseconds, that this result may be treated as fresh (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ @@ -1521,15 +1931,29 @@ "format": "const", "const": "roots/list" }, - "ListTasksResult": { + "ListToolsResult": { "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Scope describing who may cache this result (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "anyOf": [ + { + "$ref": "#/definitions/CacheScope" + }, + { + "type": "null" + } + ] }, "nextCursor": { "type": [ @@ -1537,31 +1961,15 @@ "null" ] }, - "tasks": { - "type": "array", - "items": { - "$ref": "#/definitions/Task" - } - } - }, - "required": [ - "tasks" - ] - }, - "ListToolsResult": { - "type": "object", - "properties": { - "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "nextCursor": { - "type": [ - "string", - "null" + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } ] }, "tools": { @@ -1569,6 +1977,15 @@ "items": { "$ref": "#/definitions/Tool" } + }, + "ttlMs": { + "description": "Time, in milliseconds, that this result may be treated as fresh (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ @@ -1600,11 +2017,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "data": { "description": "The actual log data" @@ -1631,6 +2051,11 @@ "data" ] }, + "MetaObject": { + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", + "type": "object", + "additionalProperties": true + }, "ModelHint": { "description": "A hint suggesting a preferred model name or family.\n\nModel hints are advisory suggestions that help clients choose appropriate\nmodels. They can be specific model names or general families like \"claude\" or \"gpt\".", "type": "object", @@ -1761,10 +2186,10 @@ "type": "object", "properties": { "method": { - "$ref": "#/definitions/ElicitationCompletionNotificationMethod" + "$ref": "#/definitions/SubscriptionsAcknowledgedNotificationMethod" }, "params": { - "$ref": "#/definitions/ElicitationResponseNotificationParam" + "$ref": "#/definitions/SubscriptionsAcknowledgedNotificationParams" } }, "required": [ @@ -1779,7 +2204,7 @@ "$ref": "#/definitions/TaskStatusNotificationMethod" }, "params": { - "$ref": "#/definitions/TaskStatusNotificationParam" + "$ref": "#/definitions/TaskStatusNotificationParams" } }, "required": [ @@ -1787,6 +2212,16 @@ "params" ] }, + "NotificationMetaObject": { + "description": "Metadata reserved by MCP on notifications. Extension keys are also allowed.", + "type": "object", + "properties": { + "io.modelcontextprotocol/subscriptionId": { + "$ref": "#/definitions/NumberOrString" + } + }, + "additionalProperties": true + }, "NotificationNoParam": { "type": "object", "properties": { @@ -1954,11 +2389,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "message": { "description": "An optional message describing the current progress.", @@ -2002,11 +2440,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "arguments": { "type": [ @@ -2117,11 +2558,25 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] + }, + "cacheScope": { + "description": "Scope describing who may cache this result (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "anyOf": [ + { + "$ref": "#/definitions/CacheScope" + }, + { + "type": "null" + } + ] }, "contents": { "description": "The actual content of the resource", @@ -2129,6 +2584,26 @@ "items": { "$ref": "#/definitions/ResourceContents" } + }, + "resultType": { + "description": "Result type discriminator (SEP-2322). Required by the [spec schema]\nfor servers implementing protocol version `2026-07-28`, but optional\nhere because this type also models results from older protocol\nversions, which do not carry the field: `None` means absent on the\nwire, and per the spec \"the client MUST treat the absent field as\n`\"complete\"`\". Constructors default to `Some(ResultType::COMPLETE)`;\nthe server handler clears the field when responding to peers that\nnegotiated an older version.\n\n[spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/271ecc9accafdd9b83a3c869fa67c22953b2af80/schema/2026-07-28/schema.ts#L219-L235", + "anyOf": [ + { + "$ref": "#/definitions/ResultType" + }, + { + "type": "null" + } + ] + }, + "ttlMs": { + "description": "Time, in milliseconds, that this result may be treated as fresh (SEP-2549).\nRequired by spec version 2026-07-28, but optional here to maintain compatibility\nwith older spec versions.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 } }, "required": [ @@ -2167,6 +2642,28 @@ "params" ] }, + "RequestMetaObject": { + "description": "Metadata reserved by MCP on requests. Extension keys are also allowed.", + "type": "object", + "properties": { + "io.modelcontextprotocol/clientCapabilities": { + "$ref": "#/definitions/ClientCapabilities" + }, + "io.modelcontextprotocol/clientInfo": { + "$ref": "#/definitions/Implementation" + }, + "io.modelcontextprotocol/logLevel": { + "$ref": "#/definitions/LoggingLevel" + }, + "io.modelcontextprotocol/protocolVersion": { + "type": "string" + }, + "progressToken": { + "$ref": "#/definitions/ProgressToken" + } + }, + "additionalProperties": true + }, "RequestNoParam": { "type": "object", "properties": { @@ -2195,11 +2692,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this resource.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this resource.", @@ -2273,11 +2773,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "mimeType": { "type": [ @@ -2301,11 +2804,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "blob": { "type": "string" @@ -2338,11 +2844,14 @@ "properties": { "_meta": { "description": "Optional protocol-level metadata for this resource template.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this template.", @@ -2410,11 +2919,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] }, "uri": { "description": "The URI of the resource that was updated", @@ -2442,6 +2954,10 @@ } } }, + "ResultType": { + "description": "Indicates the type of a result object, allowing the client to\ndetermine how to parse the response.\n\nThe spec defines this as an open string (`\"complete\" | \"input_required\" | string`),\nso unknown values are preserved rather than rejected. Servers implementing this\nprotocol version MUST include `resultType` in every result. For backward\ncompatibility, clients MUST treat an absent field as `\"complete\"`.\n\nOrdinary results model the field as `Option`: `None` means the\nfield is absent on the wire. Constructors default to `Some(COMPLETE)`, and\nthe server handler strips the `\"complete\"` discriminator before responding\nto peers that negotiated a protocol version older than `2026-07-28`, so\nlegacy sessions keep their historical wire shape (see\n[`ServerResult::strip_result_type_for_legacy_peer`]).", + "type": "string" + }, "Role": { "description": "Represents the role of a participant in a conversation or message exchange.\n\nUsed in sampling and chat contexts to distinguish between different\ntypes of message senders in the conversation flow.", "oneOf": [ @@ -2457,6 +2973,40 @@ } ] }, + "RootsCapabilities": { + "description": "Roots capability. Deprecated by SEP-2577; remains functional and will be\nremoved in a future release.\nSee .", + "type": "object", + "properties": { + "listChanged": { + "type": [ + "boolean", + "null" + ] + } + } + }, + "SamplingCapability": { + "description": "Sampling capability with optional sub-capabilities (SEP-1577).\n\nDeprecated by SEP-2577; remains functional and will be removed in a future\nrelease.\nSee .", + "type": "object", + "properties": { + "context": { + "description": "Support for `includeContext` (soft-deprecated)", + "type": [ + "object", + "null" + ], + "additionalProperties": true + }, + "tools": { + "description": "Support for `tools` and `toolChoice` parameters", + "type": [ + "object", + "null" + ], + "additionalProperties": true + } + } + }, "SamplingContent": { "description": "Single or array content wrapper (SEP-1577).", "anyOf": [ @@ -2476,11 +3026,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "content": { "description": "The actual content of the message (text, image, audio, tool use, or tool result)", @@ -2598,19 +3151,6 @@ } ] }, - "SamplingTaskCapability": { - "description": "Sampling task capability. Deprecated by SEP-2577; remains functional and\nwill be removed in a future release.\nSee .", - "type": "object", - "properties": { - "createMessage": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "ServerCapabilities": { "title": "Builder", "description": "```rust\n# use rmcp::model::ServerCapabilities;\nlet cap = ServerCapabilities::builder()\n .enable_experimental()\n .enable_prompts()\n .enable_resources()\n .enable_tools()\n .enable_tool_list_changed()\n .build();\n```", @@ -2672,16 +3212,6 @@ } ] }, - "tasks": { - "anyOf": [ - { - "$ref": "#/definitions/TasksCapability" - }, - { - "type": "null" - } - ] - }, "tools": { "anyOf": [ { @@ -2696,6 +3226,9 @@ }, "ServerResult": { "anyOf": [ + { + "$ref": "#/definitions/DiscoverResult" + }, { "$ref": "#/definitions/InitializeResult" }, @@ -2717,6 +3250,9 @@ { "$ref": "#/definitions/ReadResourceResult" }, + { + "$ref": "#/definitions/SubscriptionsListenResult" + }, { "$ref": "#/definitions/ListToolsResult" }, @@ -2726,20 +3262,17 @@ { "$ref": "#/definitions/CreateTaskResult" }, - { - "$ref": "#/definitions/ListTasksResult" - }, { "$ref": "#/definitions/GetTaskResult" }, - { - "$ref": "#/definitions/CancelTaskResult" - }, { "$ref": "#/definitions/CallToolResult" }, { - "$ref": "#/definitions/GetTaskPayloadResult" + "$ref": "#/definitions/InputRequiredResult" + }, + { + "$ref": "#/definitions/TaskAckResult" }, { "$ref": "#/definitions/EmptyObject" @@ -2857,138 +3390,135 @@ "format": "const", "const": "string" }, - "Task": { - "description": "Primary Task object that surfaces metadata during the task lifecycle.\n\nPer spec, `lastUpdatedAt` and `ttl` are required fields.\n`ttl` is nullable (`null` means unlimited retention).", + "SubscriptionFilter": { + "description": "Notification categories a client opts in to on a `subscriptions/listen` stream.", "type": "object", "properties": { - "createdAt": { - "description": "ISO-8601 creation timestamp.", - "type": "string" + "promptsListChanged": { + "type": "boolean" }, - "lastUpdatedAt": { - "description": "ISO-8601 timestamp for the most recent status change.", - "type": "string" + "resourceSubscriptions": { + "type": "array", + "items": { + "type": "string" + } }, - "pollInterval": { - "description": "Suggested polling interval (milliseconds).", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 + "resourcesListChanged": { + "type": "boolean" }, - "status": { - "description": "Current lifecycle status (see [`TaskStatus`]).", - "allOf": [ - { - "$ref": "#/definitions/TaskStatus" - } - ] + "toolsListChanged": { + "type": "boolean" + } + } + }, + "SubscriptionsAcknowledgedNotificationMethod": { + "type": "string", + "format": "const", + "const": "notifications/subscriptions/acknowledged" + }, + "SubscriptionsAcknowledgedNotificationParams": { + "description": "Parameters reporting the accepted subset of a subscription filter.", + "type": "object", + "properties": { + "_meta": { + "$ref": "#/definitions/NotificationMetaObject" }, - "statusMessage": { - "description": "Optional human-readable status message for UI surfaces.", - "type": [ - "string", - "null" - ] - }, - "taskId": { - "description": "Unique task identifier generated by the receiver.", - "type": "string" - }, - "ttl": { - "description": "Retention window in milliseconds that the receiver agreed to honor.\n`None` (serialized as `null`) means unlimited retention.", - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 + "notifications": { + "$ref": "#/definitions/SubscriptionFilter" } }, "required": [ - "taskId", - "status", - "createdAt", - "lastUpdatedAt" + "notifications" ] }, - "TaskMetadata": { - "description": "Metadata for augmenting a request with task execution (spec `TaskMetadata`).", + "SubscriptionsListenResult": { + "description": "Final response indicating that a subscription ended gracefully.", "type": "object", "properties": { - "ttl": { - "type": [ - "integer", - "null" - ], - "format": "uint64", - "minimum": 0 + "_meta": { + "$ref": "#/definitions/SubscriptionsListenResultMeta" + }, + "resultType": { + "$ref": "#/definitions/ResultType" } - } + }, + "required": [ + "resultType", + "_meta" + ] }, - "TaskRequestsCapability": { - "description": "Request types that support task-augmented execution.", + "SubscriptionsListenResultMeta": { "type": "object", "properties": { - "elicitation": { + "io.modelcontextprotocol/serverInfo": { + "description": "Identifies the server software producing the response. Servers SHOULD include this field on every response unless specifically configured not to do so.", + "allOf": [ + { + "$ref": "#/definitions/Implementation" + } + ] + }, + "io.modelcontextprotocol/subscriptionId": { + "$ref": "#/definitions/NumberOrString" + } + }, + "additionalProperties": true, + "required": [ + "io.modelcontextprotocol/subscriptionId" + ] + }, + "TaskAckResult": { + "description": "Empty acknowledgement for `tasks/update` and `tasks/cancel` (SEP-2663).\n\nThe spec requires these acks to be empty results carrying the SEP-2322\n`resultType: \"complete\"` discriminator; task state changes are observed\nvia the next `tasks/get`.", + "type": "object", + "properties": { + "_meta": { "anyOf": [ { - "$ref": "#/definitions/ElicitationTaskCapability" + "$ref": "#/definitions/MetaObject" }, { "type": "null" } ] }, - "sampling": { - "anyOf": [ + "resultType": { + "description": "Always `\"complete\"`.", + "allOf": [ { - "$ref": "#/definitions/SamplingTaskCapability" - }, - { - "type": "null" - } - ] - }, - "tools": { - "anyOf": [ - { - "$ref": "#/definitions/ToolsTaskCapability" - }, - { - "type": "null" + "$ref": "#/definitions/ResultType" } ] } - } + }, + "required": [ + "resultType" + ] }, "TaskStatus": { - "description": "Canonical task lifecycle status as defined by SEP-1686.", + "description": "Canonical task lifecycle status (SEP-2663).", "oneOf": [ { - "description": "The receiver accepted the request and is currently working on it.", + "description": "The request is currently being processed.", "type": "string", "const": "working" }, { - "description": "The receiver requires additional input before work can continue.", + "description": "The server needs input from the client before the task can proceed.", "type": "string", "const": "input_required" }, { - "description": "The underlying operation completed successfully and the result is ready.", + "description": "The request completed successfully and the result is available.\nThis includes tool calls that returned results with `isError: true`.", "type": "string", "const": "completed" }, { - "description": "The underlying operation failed and will not continue.", + "description": "The request failed due to a JSON-RPC error during execution.", "type": "string", "const": "failed" }, { - "description": "The task was cancelled and will not continue processing.", + "description": "The request was cancelled before completion.", "type": "string", "const": "cancelled" } @@ -2997,29 +3527,48 @@ "TaskStatusNotificationMethod": { "type": "string", "format": "const", - "const": "notifications/tasks/status" + "const": "notifications/tasks" }, - "TaskStatusNotificationParam": { - "description": "Parameters for a task status notification (spec `TaskStatusNotificationParams`).\n\nThe task fields are flattened at the top level: `NotificationParams & Task`.", + "TaskStatusNotificationParams": { + "description": "Parameters for a task status notification (spec `TaskStatusNotificationParams`).\n\nCarries a complete [`DetailedTask`] for the current status, identical to\nwhat `tasks/get` would have returned at that moment. The task fields are\nflattened at the top level: `NotificationParams & Task`.", "type": "object", "properties": { "_meta": { + "anyOf": [ + { + "$ref": "#/definitions/NotificationMetaObject" + }, + { + "type": "null" + } + ] + }, + "createdAt": { + "description": "ISO 8601 timestamp when the task was created.", + "type": "string" + }, + "error": { "type": [ "object", "null" ], "additionalProperties": true }, - "createdAt": { - "description": "ISO-8601 creation timestamp.", - "type": "string" + "inputRequests": { + "type": [ + "object", + "null" + ], + "additionalProperties": { + "$ref": "#/definitions/InputRequest" + } }, "lastUpdatedAt": { - "description": "ISO-8601 timestamp for the most recent status change.", + "description": "ISO 8601 timestamp when the task was last updated.", "type": "string" }, - "pollInterval": { - "description": "Suggested polling interval (milliseconds).", + "pollIntervalMs": { + "description": "Suggested polling interval in integer milliseconds. Clients SHOULD honor\nthis value to avoid overwhelming the server. This value MAY change over\nthe lifetime of a task.", "type": [ "integer", "null" @@ -3027,8 +3576,15 @@ "format": "uint64", "minimum": 0 }, + "result": { + "type": [ + "object", + "null" + ], + "additionalProperties": true + }, "status": { - "description": "Current lifecycle status (see [`TaskStatus`]).", + "description": "Current task status.", "allOf": [ { "$ref": "#/definitions/TaskStatus" @@ -3036,18 +3592,18 @@ ] }, "statusMessage": { - "description": "Optional human-readable status message for UI surfaces.", + "description": "Optional message describing the current task state.\nThis MAY be exposed to the end-user or model.", "type": [ "string", "null" ] }, "taskId": { - "description": "Unique task identifier generated by the receiver.", + "description": "Stable identifier for this task, generated by the server.", "type": "string" }, - "ttl": { - "description": "Retention window in milliseconds that the receiver agreed to honor.\n`None` (serialized as `null`) means unlimited retention.", + "ttlMs": { + "description": "Time-to-live duration from creation in integer milliseconds; `None`\n(serialized as `null`) means unlimited. The server may discard the task\nafter the TTL elapses. This value MAY change over the lifetime of a task.", "type": [ "integer", "null" @@ -3063,67 +3619,20 @@ "lastUpdatedAt" ] }, - "TaskSupport": { - "description": "Per-tool task support mode as defined in the MCP specification.\n\nThis enum indicates whether a tool supports task-based invocation,\nallowing clients to know how to properly call the tool.\n\nSee [Tool-Level Negotiation](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks#tool-level-negotiation).", - "oneOf": [ - { - "description": "Clients MUST NOT invoke this tool as a task (default behavior).", - "type": "string", - "const": "forbidden" - }, - { - "description": "Clients MAY invoke this tool as either a task or a normal call.", - "type": "string", - "const": "optional" - }, - { - "description": "Clients MUST invoke this tool as a task.", - "type": "string", - "const": "required" - } - ] - }, - "TasksCapability": { - "description": "Task capabilities shared by client and server.", - "type": "object", - "properties": { - "cancel": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "list": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, - "requests": { - "anyOf": [ - { - "$ref": "#/definitions/TaskRequestsCapability" - }, - { - "type": "null" - } - ] - } - } - }, "TextContent": { "description": "Text content block (spec `TextContent`).", "type": "object", "properties": { "_meta": { "description": "Optional protocol-level metadata for this content block.", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional annotations describing how the client should use this content.", @@ -3256,11 +3765,14 @@ "properties": { "_meta": { "description": "Optional additional metadata for this tool", - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "annotations": { "description": "Optional additional tool information.", @@ -3280,17 +3792,6 @@ "null" ] }, - "execution": { - "description": "Execution-related configuration including task support mode.", - "anyOf": [ - { - "$ref": "#/definitions/ToolExecution" - }, - { - "type": "null" - } - ] - }, "icons": { "description": "Optional list of icons for the tool", "type": [ @@ -3409,23 +3910,6 @@ } ] }, - "ToolExecution": { - "description": "Execution-related configuration for a tool.\n\nThis struct contains settings that control how a tool should be executed,\nincluding task support configuration.", - "type": "object", - "properties": { - "taskSupport": { - "description": "Indicates whether this tool supports task-based invocation.\n\nWhen not present or set to `Forbidden`, clients MUST NOT invoke this tool as a task.\nWhen set to `Optional`, clients MAY invoke this tool as a task or normal call.\nWhen set to `Required`, clients MUST invoke this tool as a task.", - "anyOf": [ - { - "$ref": "#/definitions/TaskSupport" - }, - { - "type": "null" - } - ] - } - } - }, "ToolListChangedNotificationMethod": { "type": "string", "format": "const", @@ -3436,11 +3920,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "content": { "type": "array", @@ -3454,13 +3941,7 @@ "null" ] }, - "structuredContent": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - }, + "structuredContent": true, "toolUseId": { "type": "string" } @@ -3476,11 +3957,14 @@ "type": "object", "properties": { "_meta": { - "type": [ - "object", - "null" - ], - "additionalProperties": true + "anyOf": [ + { + "$ref": "#/definitions/MetaObject" + }, + { + "type": "null" + } + ] }, "id": { "type": "string" @@ -3511,18 +3995,6 @@ } } }, - "ToolsTaskCapability": { - "type": "object", - "properties": { - "call": { - "type": [ - "object", - "null" - ], - "additionalProperties": true - } - } - }, "UntitledItems": { "description": "Items for untitled multi-select options", "type": "object", @@ -3632,6 +4104,10 @@ "type", "enum" ] + }, + "UrlElicitationCapability": { + "description": "Capability for URL mode elicitation.", + "type": "object" } } } \ No newline at end of file diff --git a/crates/rmcp/tests/test_meta_helpers.rs b/crates/rmcp/tests/test_meta_helpers.rs index a22a420..5f9128a 100644 --- a/crates/rmcp/tests/test_meta_helpers.rs +++ b/crates/rmcp/tests/test_meta_helpers.rs @@ -1,6 +1,8 @@ #![allow(deprecated)] -use rmcp::model::{ClientCapabilities, Implementation, LoggingLevel, Meta, ProtocolVersion}; +use rmcp::model::{ + ClientCapabilities, Implementation, LoggingLevel, ProtocolVersion, RequestMetaObject, +}; use serde_json::json; const META_KEY_PROTOCOL_VERSION: &str = "io.modelcontextprotocol/protocolVersion"; @@ -10,7 +12,7 @@ const META_KEY_LOG_LEVEL: &str = "io.modelcontextprotocol/logLevel"; #[test] fn meta_setters_store_sep_2575_values() { - let mut meta = Meta::new(); + let mut meta = RequestMetaObject::new(); meta.set_protocol_version(ProtocolVersion::V_2026_07_28); meta.set_client_info(Implementation::new("test-client", "1.0.0")); meta.set_client_capabilities(ClientCapabilities::default()); @@ -28,9 +30,32 @@ fn meta_setters_store_sep_2575_values() { assert_eq!(meta.get(META_KEY_LOG_LEVEL), Some(&json!("warning"))); } +#[test] +fn with_client_context_sets_all_required_sep_2575_fields() { + let meta = RequestMetaObject::with_client_context( + ProtocolVersion::V_2026_07_28, + Implementation::new("test-client", "1.0.0"), + ClientCapabilities::default(), + ); + + assert!( + meta.missing_required_keys(&ProtocolVersion::V_2026_07_28) + .is_empty() + ); + assert_eq!(meta.protocol_version(), Some(ProtocolVersion::V_2026_07_28)); + assert_eq!( + meta.client_info(), + Some(Implementation::new("test-client", "1.0.0")) + ); + assert_eq!( + meta.client_capabilities(), + Some(ClientCapabilities::default()) + ); +} + #[test] fn meta_accessors_decode_wire_values() { - let meta: Meta = serde_json::from_value(json!({ + let meta: RequestMetaObject = serde_json::from_value(json!({ "progressToken": "progress-1", "io.modelcontextprotocol/protocolVersion": "2026-07-28", "io.modelcontextprotocol/clientInfo": { @@ -58,7 +83,7 @@ fn meta_accessors_decode_wire_values() { #[test] fn meta_accessors_ignore_missing_or_malformed_values() { - let meta: Meta = serde_json::from_value(json!({ + let meta: RequestMetaObject = serde_json::from_value(json!({ "io.modelcontextprotocol/protocolVersion": 20260728, "io.modelcontextprotocol/clientInfo": "not an implementation", "io.modelcontextprotocol/clientCapabilities": "not capabilities", diff --git a/crates/rmcp/tests/test_mrtr_behavior.rs b/crates/rmcp/tests/test_mrtr_behavior.rs new file mode 100644 index 0000000..d9df5af --- /dev/null +++ b/crates/rmcp/tests/test_mrtr_behavior.rs @@ -0,0 +1,672 @@ +//! Behavior and edge-case coverage for SEP-2322 multi round-trip requests (MRTR). +//! +//! These tests drive a real client/server pair over an in-memory duplex stream +//! and exercise the auto fulfill/retry loop, the manual `*_once` escape hatch, +//! and the server-side version gating. + +// Sampling/Roots are SEP-2577-deprecated but still used to model MRTR input requests. +#![allow(deprecated)] +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use rmcp::{ + ClientHandler, ServerHandler, + handler::server::{ + tool::{InputResponses as ToolInputResponses, RequestState}, + wrapper::Parameters, + }, + model::*, + service::{RequestContext, RoleClient, RoleServer, Service, ServiceError, serve_directly}, + tool, tool_handler, tool_router, +}; +use schemars::JsonSchema; +use serde::Deserialize; +use serde_json::json; + +/// A `requestState` value with characters that must survive a byte-exact echo: +/// dots (the codec delimiter), base64 punctuation, whitespace, and quotes. +const TRICKY_STATE: &str = "st.ate/with+special=chars and spaces \"quotes\"\n\ttab"; + +// ============================================================================= +// Test handlers +// ============================================================================= + +/// A stateless MRTR server whose behavior is selected by the tool/prompt/resource +/// name. Round progression is derived entirely from `request_state` and +/// `input_responses`, as required by the stateless MRTR pattern. +#[derive(Clone, Default)] +struct MrtrServer { + calls: Arc, +} + +fn elicitation_request(message: &str) -> InputRequest { + InputRequest::Elicitation(ElicitRequest::new( + ElicitRequestParams::FormElicitationParams { + meta: None, + message: message.into(), + requested_schema: serde_json::from_value(json!({ + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"] + })) + .unwrap(), + }, + )) +} + +fn sampling_request() -> InputRequest { + InputRequest::CreateMessage(CreateMessageRequest::new(CreateMessageRequestParams::new( + vec![SamplingMessage::user_text("What is the capital of France?")], + 100, + ))) +} + +fn roots_request() -> InputRequest { + InputRequest::ListRoots(ListRootsRequest::default()) +} + +fn single_elicitation(state: &str) -> InputRequiredResult { + let mut requests = InputRequests::new(); + requests.insert("answer".to_string(), elicitation_request("Name?")); + InputRequiredResult::new(Some(requests), Some(state.into())) +} + +#[derive(Clone)] +struct MacroMrtrServer; + +#[derive(Deserialize, JsonSchema)] +struct MacroMrtrArguments { + greeting: String, +} + +#[tool_router] +impl MacroMrtrServer { + #[tool(description = "Greet a user after collecting their name")] + async fn greet( + &self, + Parameters(arguments): Parameters, + RequestState(request_state): RequestState, + ToolInputResponses(input_responses): ToolInputResponses, + ) -> Result { + match request_state.as_deref() { + None => Ok(single_elicitation("macro-state").into()), + Some("macro-state") => { + let name = input_responses + .as_ref() + .and_then(|responses| responses.get("answer")) + .and_then(|response| response["content"]["name"].as_str()) + .ok_or_else(|| ErrorData::invalid_params("missing name response", None))?; + Ok(CallToolResult::success(vec![ContentBlock::text(format!( + "{}, {name}", + arguments.greeting + ))]) + .into()) + } + Some(other) => Err(ErrorData::invalid_params( + format!("unexpected request state {other:?}"), + None, + )), + } + } +} + +#[tool_handler] +impl ServerHandler for MacroMrtrServer { + fn get_info(&self) -> ServerConfig { + let mut info = ServerConfig::new(ServerCapabilities::builder().enable_tools().build()); + info.protocol_version = ProtocolVersion::V_2026_07_28; + info + } +} + +impl MrtrServer { + fn call_tool_impl( + &self, + request: CallToolRequestParams, + ) -> Result { + let responses = request.input_responses.as_ref(); + let state = request.request_state.as_deref(); + match request.name.as_ref() { + // Single round: one elicitation, then complete. + "single" => match responses { + None => Ok(single_elicitation("state-single").into()), + Some(map) => { + if state != Some("state-single") { + return Err(ErrorData::internal_error("request_state not echoed", None)); + } + let answer = map + .get("answer") + .ok_or_else(|| ErrorData::internal_error("missing answer", None))?; + if answer["action"] != "accept" || answer["content"]["name"] != "Ferris" { + return Err(ErrorData::internal_error("unexpected elicit result", None)); + } + Ok(CallToolResult::success(vec![ContentBlock::text("done")]).into()) + } + }, + // Two elicitation rounds before completing. + "multi_round" => match state { + None => Ok(single_elicitation("round-1").into()), + Some("round-1") => Ok(single_elicitation("round-2").into()), + Some("round-2") => { + Ok(CallToolResult::success(vec![ContentBlock::text("multi-done")]).into()) + } + Some(other) => Err(ErrorData::internal_error( + format!("unexpected round state {other:?}"), + None, + )), + }, + // Several input requests fulfilled concurrently in a single round. + "multi_request" => match responses { + None => { + let mut requests = InputRequests::new(); + requests.insert("form".to_string(), elicitation_request("Name?")); + requests.insert("sample".to_string(), sampling_request()); + requests.insert("roots".to_string(), roots_request()); + Ok(InputRequiredResult::new(Some(requests), Some("multi-req".into())).into()) + } + Some(map) => { + for key in ["form", "sample", "roots"] { + if !map.contains_key(key) { + return Err(ErrorData::internal_error( + format!("missing response for {key}"), + None, + )); + } + } + Ok(CallToolResult::success(vec![ContentBlock::text("multi-req-done")]).into()) + } + }, + // State-only load shedding: two state-only rounds, then complete. + "state_only" => match state { + None => Ok(InputRequiredResult::from_request_state("so-1").into()), + Some("so-1") => Ok(InputRequiredResult::from_request_state("so-2").into()), + Some("so-2") => { + Ok(CallToolResult::success(vec![ContentBlock::text("state-done")]).into()) + } + Some(other) => Err(ErrorData::internal_error( + format!("unexpected state {other:?}"), + None, + )), + }, + // Never completes: used to exercise the max-rounds cap. + "loops" => Ok(single_elicitation("loop").into()), + // Triggers a failure inside the client's elicitation handler. + "handler_error" => { + let mut requests = InputRequests::new(); + requests.insert("answer".to_string(), elicitation_request("FAIL")); + Ok(InputRequiredResult::new(Some(requests), Some("state".into())).into()) + } + // Verifies the client echoes `request_state` byte-for-byte. + "echo_state" => match responses { + None => { + let mut requests = InputRequests::new(); + requests.insert("answer".to_string(), elicitation_request("Name?")); + Ok(InputRequiredResult::new(Some(requests), Some(TRICKY_STATE.into())).into()) + } + Some(_) => { + if state != Some(TRICKY_STATE) { + return Err(ErrorData::internal_error( + "request_state was not echoed byte-exact", + None, + )); + } + Ok(CallToolResult::success(vec![ContentBlock::text("echo-ok")]).into()) + } + }, + _ => Ok(CallToolResult::success(vec![ContentBlock::text("noop")]).into()), + } + } +} + +impl ServerHandler for MrtrServer { + fn get_info(&self) -> ServerConfig { + let mut info = ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_prompts() + .build(), + ); + info.protocol_version = ProtocolVersion::V_2026_07_28; + info + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + self.call_tool_impl(request) + } + + async fn get_prompt( + &self, + request: GetPromptRequestParams, + _context: RequestContext, + ) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + match request.request_state.as_deref() { + None => Ok(single_elicitation("prompt-1").into()), + Some("prompt-1") => Ok(GetPromptResult::new(vec![PromptMessage::new_text( + Role::Assistant, + "prompt-done", + )]) + .into()), + Some(other) => Err(ErrorData::internal_error( + format!("unexpected prompt state {other:?}"), + None, + )), + } + } + + async fn read_resource( + &self, + request: ReadResourceRequestParams, + _context: RequestContext, + ) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + match request.request_state.as_deref() { + None => Ok(single_elicitation("res-1").into()), + Some("res-1") => Ok(ReadResourceResult::new(vec![ResourceContents::text( + "resource-done", + request.uri, + )]) + .into()), + Some(other) => Err(ErrorData::internal_error( + format!("unexpected resource state {other:?}"), + None, + )), + } + } +} + +/// A client that fulfills every kind of MRTR input request. Elicitation fails +/// deliberately when the prompt message is `"FAIL"`. +#[derive(Clone, Default)] +struct MrtrClient; + +impl ClientHandler for MrtrClient { + async fn create_elicitation( + &self, + request: ElicitRequestParams, + _context: RequestContext, + ) -> Result { + if let ElicitRequestParams::FormElicitationParams { message, .. } = &request + && message == "FAIL" + { + return Err(ErrorData::internal_error( + "elicitation handler failed", + None, + )); + } + Ok(ElicitResult::new(ElicitationAction::Accept).with_content(json!({ "name": "Ferris" }))) + } + + async fn create_message( + &self, + _request: CreateMessageRequestParams, + _context: RequestContext, + ) -> Result { + Ok(CreateMessageResult::new( + SamplingMessage::assistant_text("Paris."), + "test-model".into(), + ) + .with_stop_reason(CreateMessageResult::STOP_REASON_END_TURN)) + } + + async fn list_roots( + &self, + _context: RequestContext, + ) -> Result { + Ok(ListRootsResult::new(vec![Root::new("file:///workspace")])) + } +} + +// ============================================================================= +// Harness +// ============================================================================= + +fn client_info(protocol_version: ProtocolVersion) -> InitializeRequestParams { + InitializeRequestParams::new( + ClientCapabilities::builder().enable_elicitation().build(), + Implementation::new("mrtr-test-client", "0.0.0"), + ) + .with_protocol_version(protocol_version) +} + +fn server_info(protocol_version: ProtocolVersion) -> InitializeResult { + let mut info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build()); + info.protocol_version = protocol_version; + info +} + +/// Runs `body` inside a `LocalSet` so `spawn_local` (used when the `local` +/// feature is active) is available, wiring up a connected client/server pair. +async fn with_pair( + server: S, + client_protocol: ProtocolVersion, + body: F, +) -> anyhow::Result<()> +where + S: Service, + F: FnOnce(rmcp::service::RunningService) -> Fut, + Fut: std::future::Future>, +{ + tokio::task::LocalSet::new() + .run_until(async move { + let (server_transport, client_transport) = tokio::io::duplex(8192); + let server_peer_info = client_info(client_protocol); + let client_peer_info = server_info(server_peer_info.protocol_version.clone()); + let server_task = tokio::task::spawn_local(async move { + let running = serve_directly::( + server, + server_transport, + Some(server_peer_info), + ); + running.waiting().await?; + anyhow::Ok(()) + }); + + let client = serve_directly::( + MrtrClient, + client_transport, + Some(client_peer_info.into()), + ); + + let result = body(client).await; + + server_task.abort(); + result + }) + .await +} + +// ============================================================================= +// Tests +// ============================================================================= + +#[tokio::test(flavor = "current_thread")] +async fn client_auto_fulfills_input_required_tool_call() -> anyhow::Result<()> { + let server = MrtrServer::default(); + let calls = server.calls.clone(); + with_pair(server, ProtocolVersion::V_2026_07_28, |client| async move { + let result = client + .call_tool(CallToolRequestParams::new("single")) + .await?; + assert_eq!(result.content[0].as_text().unwrap().text, "done"); + assert_eq!(calls.load(Ordering::SeqCst), 2); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn tool_macro_receives_mrtr_retry_fields() -> anyhow::Result<()> { + with_pair( + MacroMrtrServer, + ProtocolVersion::V_2026_07_28, + |client| async move { + let arguments = serde_json::from_value(json!({ "greeting": "hello" })).unwrap(); + let result = client + .call_tool(CallToolRequestParams::new("greet").with_arguments(arguments)) + .await?; + assert_eq!(result.content[0].as_text().unwrap().text, "hello, Ferris"); + Ok(()) + }, + ) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn manual_once_returns_input_required_without_retry() -> anyhow::Result<()> { + let server = MrtrServer::default(); + let calls = server.calls.clone(); + with_pair(server, ProtocolVersion::V_2026_07_28, |client| async move { + let result = client + .call_tool_once(CallToolRequestParams::new("single")) + .await?; + assert!(matches!(result, CallToolResponse::InputRequired(_))); + // A manual round makes exactly one server call and never retries. + assert_eq!(calls.load(Ordering::SeqCst), 1); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn multi_round_input_required_completes() -> anyhow::Result<()> { + let server = MrtrServer::default(); + let calls = server.calls.clone(); + with_pair(server, ProtocolVersion::V_2026_07_28, |client| async move { + let result = client + .call_tool(CallToolRequestParams::new("multi_round")) + .await?; + assert_eq!(result.content[0].as_text().unwrap().text, "multi-done"); + // round 0 + two retries = 3 server calls. + assert_eq!(calls.load(Ordering::SeqCst), 3); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn multiple_input_requests_fulfilled_in_one_round() -> anyhow::Result<()> { + let server = MrtrServer::default(); + let calls = server.calls.clone(); + with_pair(server, ProtocolVersion::V_2026_07_28, |client| async move { + let result = client + .call_tool(CallToolRequestParams::new("multi_request")) + .await?; + assert_eq!(result.content[0].as_text().unwrap().text, "multi-req-done"); + assert_eq!(calls.load(Ordering::SeqCst), 2); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn state_only_input_required_completes() -> anyhow::Result<()> { + let server = MrtrServer::default(); + let calls = server.calls.clone(); + with_pair(server, ProtocolVersion::V_2026_07_28, |client| async move { + let result = client + .call_tool(CallToolRequestParams::new("state_only")) + .await?; + assert_eq!(result.content[0].as_text().unwrap().text, "state-done"); + // round 0 + two state-only retries = 3 server calls. + assert_eq!(calls.load(Ordering::SeqCst), 3); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn max_rounds_exceeded_returns_error() -> anyhow::Result<()> { + let server = MrtrServer::default(); + let calls = server.calls.clone(); + with_pair(server, ProtocolVersion::V_2026_07_28, |client| async move { + let err = client + .call_tool_with_mrtr_max_rounds(CallToolRequestParams::new("loops"), 3) + .await + .expect_err("a tool that never completes must exhaust the round cap"); + assert!(matches!( + err, + ServiceError::InputRequiredRoundsExceeded { max_rounds: 3 } + )); + assert_eq!(calls.load(Ordering::SeqCst), 3); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn client_handler_error_propagates() -> anyhow::Result<()> { + let server = MrtrServer::default(); + with_pair(server, ProtocolVersion::V_2026_07_28, |client| async move { + let err = client + .call_tool(CallToolRequestParams::new("handler_error")) + .await + .expect_err("a failing input handler must fail the whole call"); + assert!(matches!(err, ServiceError::McpError(_))); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn request_state_is_echoed_byte_exact() -> anyhow::Result<()> { + let server = MrtrServer::default(); + with_pair(server, ProtocolVersion::V_2026_07_28, |client| async move { + // The server returns an error result unless it sees TRICKY_STATE echoed + // back unchanged, so a successful completion proves the byte-exact echo. + let result = client + .call_tool(CallToolRequestParams::new("echo_state")) + .await?; + assert_eq!(result.content[0].as_text().unwrap().text, "echo-ok"); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn get_prompt_auto_fulfills_input_required() -> anyhow::Result<()> { + let server = MrtrServer::default(); + with_pair(server, ProtocolVersion::V_2026_07_28, |client| async move { + let result = client.get_prompt(GetPromptRequestParams::new("p")).await?; + assert_eq!( + result.messages[0].content.as_text().unwrap().text, + "prompt-done" + ); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn read_resource_auto_fulfills_input_required() -> anyhow::Result<()> { + let server = MrtrServer::default(); + with_pair(server, ProtocolVersion::V_2026_07_28, |client| async move { + let result = client + .read_resource(ReadResourceRequestParams::new("res://x")) + .await?; + let text = match &result.contents[0] { + ResourceContents::TextResourceContents { text, .. } => text.clone(), + _ => panic!("expected text resource"), + }; + assert_eq!(text, "resource-done"); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn old_protocol_rejects_input_required() -> anyhow::Result<()> { + let server = MrtrServer::default(); + // The client negotiated 2025-11-25, so the server must refuse to emit an + // InputRequiredResult and return a protocol error instead. + with_pair(server, ProtocolVersion::V_2025_11_25, |client| async move { + let err = client + .call_tool_once(CallToolRequestParams::new("single")) + .await + .expect_err("MRTR must be rejected for pre-2026 peers"); + match err { + ServiceError::McpError(error) => { + assert!( + error.message.contains("2026-07-28"), + "unexpected error message: {}", + error.message + ); + } + other => panic!("expected an McpError, got {other:?}"), + } + Ok(()) + }) + .await +} + +#[cfg(feature = "request-state")] +#[tokio::test(flavor = "current_thread")] +async fn request_state_codec_seals_and_verifies_through_the_loop() -> anyhow::Result<()> { + use std::sync::OnceLock; + + use rmcp::model::RequestStateCodec; + + // A shared per-process signing key, mirroring how a real server would derive one. + const KEY: &[u8] = b"integration-signing-key-32-bytes!"; + const _: () = assert!( + KEY.len() >= RequestStateCodec::MIN_KEY_LENGTH, + "KEY is shorter than RequestStateCodec::MIN_KEY_LENGTH", + ); + + fn codec() -> &'static RequestStateCodec { + static CODEC: OnceLock = OnceLock::new(); + CODEC.get_or_init(|| RequestStateCodec::new_unchecked(KEY)) + } + + #[derive(Clone, Default)] + struct SealingServer; + + impl ServerHandler for SealingServer { + fn get_info(&self) -> ServerConfig { + let mut info = ServerConfig::new(ServerCapabilities::builder().enable_tools().build()); + info.protocol_version = ProtocolVersion::V_2026_07_28; + info + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + match request.request_state { + None => { + let sealed = codec() + .seal_json(&json!({ "step": 1, "tool": request.name })) + .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; + let mut requests = InputRequests::new(); + requests.insert("answer".to_string(), elicitation_request("Name?")); + Ok(InputRequiredResult::new(Some(requests), Some(sealed)).into()) + } + Some(sealed) => { + // The echoed state is untrusted; verify it before use. + let state: serde_json::Value = codec() + .open_json(&sealed) + .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; + assert_eq!(state["step"], 1); + Ok(CallToolResult::success(vec![ContentBlock::text("sealed-done")]).into()) + } + } + } + } + + tokio::task::LocalSet::new() + .run_until(async move { + let (server_transport, client_transport) = tokio::io::duplex(8192); + let server_task = tokio::task::spawn_local(async move { + let running = serve_directly::( + SealingServer, + server_transport, + Some(client_info(ProtocolVersion::V_2026_07_28)), + ); + running.waiting().await?; + anyhow::Ok(()) + }); + + let client = serve_directly::( + MrtrClient, + client_transport, + Some(server_info(ProtocolVersion::V_2026_07_28).into()), + ); + + let result = client + .call_tool(CallToolRequestParams::new("sealed")) + .await?; + assert_eq!(result.content[0].as_text().unwrap().text, "sealed-done"); + + server_task.abort(); + anyhow::Ok(()) + }) + .await +} diff --git a/crates/rmcp/tests/test_notification.rs b/crates/rmcp/tests/test_notification.rs index 073396e..dafb0df 100644 --- a/crates/rmcp/tests/test_notification.rs +++ b/crates/rmcp/tests/test_notification.rs @@ -1,11 +1,12 @@ #![cfg(not(feature = "local"))] +#![allow(deprecated)] use std::sync::Arc; use rmcp::{ ClientHandler, ServerHandler, ServiceExt, model::{ ClientNotification, CustomNotification, ResourceUpdatedNotificationParam, - ServerCapabilities, ServerInfo, ServerNotification, SubscribeRequestParams, + ServerCapabilities, ServerConfig, ServerNotification, SubscribeRequestParams, }, }; use serde_json::json; @@ -15,8 +16,8 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; struct Server {} impl ServerHandler for Server { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_resources() .enable_resources_subscribe() diff --git a/crates/rmcp/tests/test_progress_subscriber.rs b/crates/rmcp/tests/test_progress_subscriber.rs index 18f91c2..8df1b7c 100644 --- a/crates/rmcp/tests/test_progress_subscriber.rs +++ b/crates/rmcp/tests/test_progress_subscriber.rs @@ -3,7 +3,9 @@ use futures::StreamExt; use rmcp::{ ClientHandler, Peer, RoleServer, ServerHandler, ServiceExt, handler::{client::progress::ProgressDispatcher, server::tool::ToolRouter}, - model::{CallToolRequestParams, ClientRequest, Meta, ProgressNotificationParam, Request}, + model::{ + CallToolRequestParams, ClientRequest, ProgressNotificationParam, Request, RequestMetaObject, + }, service::PeerRequestOptions, tool, tool_handler, tool_router, }; @@ -61,7 +63,7 @@ impl Default for MyServer { impl MyServer { #[tool] pub async fn some_progress( - meta: Meta, + meta: RequestMetaObject, client: Peer, ) -> Result<(), rmcp::ErrorData> { let progress_token = meta diff --git a/crates/rmcp/tests/test_prompt_macros.rs b/crates/rmcp/tests/test_prompt_macros.rs index 7a00249..e6d65b6 100644 --- a/crates/rmcp/tests/test_prompt_macros.rs +++ b/crates/rmcp/tests/test_prompt_macros.rs @@ -4,14 +4,12 @@ use std::sync::Arc; use rmcp::{ - ClientHandler, RoleServer, ServerHandler, ServiceExt, + ClientHandler, ServerHandler, ServiceExt, handler::server::{router::prompt::PromptRouter, wrapper::Parameters}, model::{ - ClientInfo, ContentBlock, GetPromptRequestParams, GetPromptResult, ListPromptsResult, - PaginatedRequestParams, PromptMessage, Role, + ClientConfig, ContentBlock, GetPromptRequestParams, GetPromptResult, PromptMessage, Role, }, prompt, prompt_handler, prompt_router, - service::RequestContext, }; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; @@ -299,8 +297,8 @@ fn test_optional_field_schema_generation_via_macro() { struct DummyClientHandler {} impl ClientHandler for DummyClientHandler { - fn get_info(&self) -> ClientInfo { - ClientInfo::default() + fn get_info(&self) -> ClientConfig { + ClientConfig::default() } } diff --git a/crates/rmcp/tests/test_prompt_routers.rs b/crates/rmcp/tests/test_prompt_routers.rs index eecdb1e..e730579 100644 --- a/crates/rmcp/tests/test_prompt_routers.rs +++ b/crates/rmcp/tests/test_prompt_routers.rs @@ -20,12 +20,6 @@ struct Request { fields: HashMap, } -#[derive(Debug, schemars::JsonSchema, serde::Deserialize, serde::Serialize)] -struct Sum { - a: i32, - b: i32, -} - #[rmcp::prompt_router(router = "test_router")] impl TestHandler { #[rmcp::prompt] diff --git a/crates/rmcp/tests/test_protocol_version_negotiation.rs b/crates/rmcp/tests/test_protocol_version_negotiation.rs index 44a314e..1453e0f 100644 --- a/crates/rmcp/tests/test_protocol_version_negotiation.rs +++ b/crates/rmcp/tests/test_protocol_version_negotiation.rs @@ -1,20 +1,98 @@ //! Tests for protocol version negotiation in the default ServerHandler::initialize impl. //! -//! Known versions are echoed back; unknown versions fall back to LATEST. +//! Handshake versions are echoed back; every other version falls back to one +//! the server can serve over `initialize`. #![cfg(not(feature = "local"))] #![cfg(feature = "client")] +use std::{ + borrow::Cow, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; + use rmcp::{ - ClientHandler, ServerHandler, ServiceExt, - model::{ClientInfo, ProtocolVersion, ServerInfo}, + ClientHandler, ErrorData, RoleServer, ServerHandler, ServiceExt, + model::{ + ClientCapabilities, ClientConfig, ErrorCode, Implementation, InitializeRequestParams, + InitializeResult, ProtocolVersion, ServerConfig, + }, + service::{ClientInitializeError, RequestContext}, }; #[derive(Debug, Clone, Default)] struct EchoServer; impl ServerHandler for EchoServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::default() + fn get_info(&self) -> ServerConfig { + ServerConfig::default() + } +} + +/// Every known version whose lifecycle still runs the `initialize` handshake. +/// `2026-07-28` replaced the handshake with per-request metadata, so this is +/// also the list a server that has not implemented that revision supports. +const HANDSHAKE_VERSIONS: &[ProtocolVersion] = &[ + ProtocolVersion::V_2024_11_05, + ProtocolVersion::V_2025_03_26, + ProtocolVersion::V_2025_06_18, + ProtocolVersion::V_2025_11_25, +]; + +#[derive(Debug, Clone, Default)] +struct NarrowedServer; + +impl ServerHandler for NarrowedServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::default() + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(HANDSHAKE_VERSIONS) + } +} + +/// Supports only revisions that have no `initialize` handshake at all. +#[derive(Debug, Clone, Default)] +struct ModernOnlyServer; + +const MODERN_ONLY_VERSIONS: &[ProtocolVersion] = &[ProtocolVersion::V_2026_07_28]; + +impl ServerHandler for ModernOnlyServer { + fn get_info(&self) -> ServerConfig { + let mut info = ServerConfig::default(); + info.protocol_version = ProtocolVersion::V_2026_07_28; + info + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(MODERN_ONLY_VERSIONS) + } +} + +/// Narrows the supported versions *and* overrides `initialize`, so the +/// handler's own answer never runs the default negotiation. The handshake layer +/// must still honor the narrowed list. +#[derive(Debug, Clone, Default)] +struct NarrowedOverridingServer; + +impl ServerHandler for NarrowedOverridingServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::default() + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(HANDSHAKE_VERSIONS) + } + + async fn initialize( + &self, + _request: InitializeRequestParams, + _context: RequestContext, + ) -> Result { + Ok(self.get_info()) } } @@ -24,31 +102,43 @@ struct VersionedClient { } impl ClientHandler for VersionedClient { - fn get_info(&self) -> ClientInfo { - let mut info = ClientInfo::default(); + fn get_info(&self) -> ClientConfig { + let mut info = ClientConfig::default(); info.protocol_version = self.protocol_version.clone(); info } } async fn negotiated_version(client_version: ProtocolVersion) -> ProtocolVersion { + negotiated_version_with(EchoServer, client_version).await +} + +async fn negotiated_version_with( + server: S, + client_version: ProtocolVersion, +) -> ProtocolVersion { + negotiate_with(server, client_version) + .await + .expect("client should connect") +} + +async fn negotiate_with( + server: S, + client_version: ProtocolVersion, +) -> Result { let (server_transport, client_transport) = tokio::io::duplex(4096); tokio::spawn(async move { - let _ = EchoServer - .serve(server_transport) - .await - .expect("server should start") - .waiting() - .await; + if let Ok(running) = server.serve(server_transport).await { + let _ = running.waiting().await; + } }); let client = VersionedClient { protocol_version: client_version, } .serve(client_transport) - .await - .expect("client should connect"); + .await?; let version = client .peer_info() @@ -57,20 +147,55 @@ async fn negotiated_version(client_version: ProtocolVersion) -> ProtocolVersion .clone(); client.cancel().await.expect("client should cancel"); - version + Ok(version) } #[tokio::test] -async fn known_version_echoed_back() { - for version in ProtocolVersion::KNOWN_VERSIONS { +async fn handshake_version_echoed_back() { + for version in HANDSHAKE_VERSIONS { let negotiated = negotiated_version(version.clone()).await; assert_eq!( negotiated, *version, - "known version {version} should be echoed back" + "handshake version {version} should be echoed back" ); } } +/// `initialize` disappeared in `2026-07-28`, so agreeing to it here would leave +/// the peers speaking a revision that has no handshake at all. +#[tokio::test] +async fn handshake_never_agrees_to_a_version_that_dropped_it() { + let negotiated = negotiated_version(ProtocolVersion::V_2026_07_28).await; + assert_eq!( + negotiated, + ProtocolVersion::LATEST, + "a version that dropped the handshake should fall back to the server's own" + ); +} + +#[tokio::test] +async fn modern_only_server_rejects_the_handshake() { + let error = negotiate_with(ModernOnlyServer, ProtocolVersion::V_2026_07_28) + .await + .expect_err("a server with no handshake version cannot answer initialize"); + let ClientInitializeError::JsonRpcError(error) = error else { + panic!("expected a JSON-RPC error, got {error:?}"); + }; + assert_eq!( + error.code, + ErrorCode::UNSUPPORTED_PROTOCOL_VERSION, + "a server with no handshake version should reject initialize" + ); + assert_eq!( + error.data, + Some(serde_json::json!({ + "requested": "2026-07-28", + "supported": ["2026-07-28"], + })), + "the rejection should name the versions the server does support" + ); +} + #[tokio::test] async fn unknown_version_falls_back_to_latest() { let unknown: ProtocolVersion = serde_json::from_str(r#""1999-01-01""#).unwrap(); @@ -81,3 +206,123 @@ async fn unknown_version_falls_back_to_latest() { "unknown version should fall back to LATEST" ); } + +#[tokio::test] +async fn narrowed_server_still_echoes_versions_it_supports() { + for version in HANDSHAKE_VERSIONS { + let negotiated = negotiated_version_with(NarrowedServer, version.clone()).await; + assert_eq!( + negotiated, *version, + "supported version {version} should be echoed back" + ); + } +} + +#[tokio::test] +async fn narrowed_server_does_not_agree_to_version_it_excludes() { + let negotiated = negotiated_version_with(NarrowedServer, ProtocolVersion::V_2026_07_28).await; + assert_eq!( + negotiated, + ProtocolVersion::V_2025_11_25, + "a version outside supported_protocol_versions should not be echoed back" + ); +} + +#[tokio::test] +async fn narrowed_server_caps_even_when_it_overrides_initialize() { + let negotiated = + negotiated_version_with(NarrowedOverridingServer, ProtocolVersion::V_2026_07_28).await; + assert_eq!( + negotiated, + ProtocolVersion::V_2025_11_25, + "the handshake layer should not raise the version above what the server supports" + ); +} + +/// Overrides `initialize` to run a side effect, then delegates the version +/// answer back to the SDK with [`ServerHandler::negotiate_initialize`]. +#[derive(Debug, Clone, Default)] +struct DelegatingServer { + initializations: Arc, +} + +impl ServerHandler for DelegatingServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::default() + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(HANDSHAKE_VERSIONS) + } + + async fn initialize( + &self, + request: InitializeRequestParams, + context: RequestContext, + ) -> Result { + self.initializations.fetch_add(1, Ordering::Relaxed); + context.peer.set_peer_info(request.clone()); + self.negotiate_initialize(&request) + } +} + +fn initialize_params(protocol_version: ProtocolVersion) -> InitializeRequestParams { + let mut params = InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("test-client", "0.0.0"), + ); + params.protocol_version = protocol_version; + params +} + +#[test] +fn negotiate_initialize_echoes_a_supported_version() { + let result = NarrowedServer + .negotiate_initialize(&initialize_params(ProtocolVersion::V_2025_06_18)) + .expect("a supported handshake version should negotiate"); + assert_eq!(result.protocol_version, ProtocolVersion::V_2025_06_18); +} + +#[test] +fn negotiate_initialize_caps_at_supported_versions() { + let result = NarrowedServer + .negotiate_initialize(&initialize_params(ProtocolVersion::V_2026_07_28)) + .expect("an unsupported version should fall back rather than fail"); + assert_eq!(result.protocol_version, ProtocolVersion::V_2025_11_25); +} + +#[test] +fn negotiate_initialize_keeps_the_rest_of_get_info() { + let server = NarrowedServer; + let result = server + .negotiate_initialize(&initialize_params(ProtocolVersion::V_2026_07_28)) + .expect("an unsupported version should fall back rather than fail"); + assert_eq!(result.capabilities, server.get_info().capabilities); +} + +#[test] +fn negotiate_initialize_rejects_when_no_handshake_version_is_supported() { + let error = ModernOnlyServer + .negotiate_initialize(&initialize_params(ProtocolVersion::V_2026_07_28)) + .expect_err("a server with no handshake version cannot answer initialize"); + assert_eq!(error.code, ErrorCode::UNSUPPORTED_PROTOCOL_VERSION); +} + +#[tokio::test] +async fn delegating_server_negotiates_like_the_default_initialize() { + let negotiated = + negotiated_version_with(DelegatingServer::default(), ProtocolVersion::V_2026_07_28).await; + assert_eq!( + negotiated, + ProtocolVersion::V_2025_11_25, + "an override that delegates should answer what the default initialize would" + ); +} + +#[tokio::test] +async fn delegating_server_still_runs_its_own_side_effect() { + let server = DelegatingServer::default(); + let initializations = Arc::clone(&server.initializations); + negotiated_version_with(server, ProtocolVersion::V_2025_06_18).await; + assert_eq!(initializations.load(Ordering::Relaxed), 1); +} diff --git a/crates/rmcp/tests/test_request_timeout_progress.rs b/crates/rmcp/tests/test_request_timeout_progress.rs index 6eeac42..f37101f 100644 --- a/crates/rmcp/tests/test_request_timeout_progress.rs +++ b/crates/rmcp/tests/test_request_timeout_progress.rs @@ -11,8 +11,8 @@ use std::{ use rmcp::{ ClientHandler, Peer, RoleServer, ServiceError, ServiceExt, model::{ - CallToolRequestParams, ClientRequest, Meta, NumberOrString, ProgressNotificationParam, - ProgressToken, Request, + CallToolRequestParams, ClientRequest, NumberOrString, ProgressNotificationParam, + ProgressToken, Request, RequestMetaObject, }, service::PeerRequestOptions, tool, tool_router, @@ -46,7 +46,7 @@ impl ProgressTimeoutServer { #[tool] async fn delayed_with_progress( &self, - meta: Meta, + meta: RequestMetaObject, client: Peer, ) -> Result<(), rmcp::ErrorData> { let progress_token = meta @@ -173,7 +173,7 @@ async fn generated_progress_token_overrides_option_meta_token() -> anyhow::Resul let client = start_pair().await?; let mut options = PeerRequestOptions::with_timeout(Duration::from_millis(75)).reset_timeout_on_progress(); - options.meta = Some(Meta::with_progress_token(ProgressToken( + options.meta = Some(RequestMetaObject::with_progress_token(ProgressToken( NumberOrString::Number(999_999), ))); diff --git a/crates/rmcp/tests/test_resource_not_found_version.rs b/crates/rmcp/tests/test_resource_not_found_version.rs index 255accb..efae45f 100644 --- a/crates/rmcp/tests/test_resource_not_found_version.rs +++ b/crates/rmcp/tests/test_resource_not_found_version.rs @@ -6,12 +6,12 @@ #![cfg(feature = "client")] use rmcp::{ - ClientHandler, RoleServer, ServerHandler, ServiceError, ServiceExt, + ClientHandler, RoleClient, RoleServer, ServerHandler, ServiceError, model::{ - ClientInfo, ErrorCode, ErrorData, ProtocolVersion, ReadResourceRequestParams, - ReadResourceResult, + ClientConfig, ErrorCode, ErrorData, InitializeResult, ProtocolVersion, + ReadResourceRequestParams, ReadResourceResponse, }, - service::RequestContext, + service::{RequestContext, serve_directly}, }; #[derive(Debug, Clone, Default)] @@ -22,7 +22,7 @@ impl ServerHandler for ResourceServer { &self, _request: ReadResourceRequestParams, _context: RequestContext, - ) -> Result { + ) -> Result { Err(ErrorData::resource_not_found("resource not found", None)) } } @@ -33,31 +33,40 @@ struct VersionedClient { } impl ClientHandler for VersionedClient { - fn get_info(&self) -> ClientInfo { - let mut info = ClientInfo::default(); + fn get_info(&self) -> ClientConfig { + let mut info = ClientConfig::default(); info.protocol_version = self.protocol_version.clone(); info } } +/// Wires the pair up directly on `client_version`. `2026-07-28` removed the +/// `initialize` handshake, so a peer on that revision is reached the way the +/// discover lifecycle leaves one: with the version already agreed. async fn not_found_code(client_version: ProtocolVersion) -> ErrorCode { let (server_transport, client_transport) = tokio::io::duplex(4096); + let client_handler = VersionedClient { + protocol_version: client_version.clone(), + }; + let mut server_peer_info = InitializeResult::default(); + server_peer_info.protocol_version = client_version; + + let server = serve_directly::( + ResourceServer, + server_transport, + Some(client_handler.get_info()), + ); let server_handle = tokio::spawn(async move { - ResourceServer - .serve(server_transport) - .await? - .waiting() - .await?; + server.waiting().await?; anyhow::Ok(()) }); - let client = VersionedClient { - protocol_version: client_version, - } - .serve(client_transport) - .await - .expect("client should connect"); + let client = serve_directly::( + client_handler, + client_transport, + Some(server_peer_info.into()), + ); let error = client .read_resource(ReadResourceRequestParams::new("missing://resource")) diff --git a/crates/rmcp/tests/test_result_type_version.rs b/crates/rmcp/tests/test_result_type_version.rs new file mode 100644 index 0000000..61240ed --- /dev/null +++ b/crates/rmcp/tests/test_result_type_version.rs @@ -0,0 +1,95 @@ +//! SEP-2322: the `resultType` discriminator follows the negotiated protocol version. +//! +//! Peers negotiating `2026-07-28` or newer receive `resultType: "complete"` on +//! ordinary results; older peers keep the legacy wire shape without the field. +#![cfg(not(feature = "local"))] +#![cfg(feature = "client")] + +use rmcp::{ + ClientHandler, RoleClient, RoleServer, ServerHandler, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ClientConfig, ContentBlock, + ErrorData, InitializeResult, ProtocolVersion, ResultType, + }, + service::{RequestContext, serve_directly}, +}; + +#[derive(Debug, Clone, Default)] +struct ToolServer; + +impl ServerHandler for ToolServer { + async fn call_tool( + &self, + _request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + Ok(CallToolResult::success(vec![ContentBlock::text("ok")]).into()) + } +} + +#[derive(Debug, Clone)] +struct VersionedClient { + protocol_version: ProtocolVersion, +} + +impl ClientHandler for VersionedClient { + fn get_info(&self) -> ClientConfig { + let mut info = ClientConfig::default(); + info.protocol_version = self.protocol_version.clone(); + info + } +} + +/// Wires the pair up directly on `client_version`. `2026-07-28` removed the +/// `initialize` handshake, so a peer on that revision is reached the way the +/// discover lifecycle leaves one: with the version already agreed. +async fn call_tool_result_type(client_version: ProtocolVersion) -> Option { + let (server_transport, client_transport) = tokio::io::duplex(4096); + + let client_handler = VersionedClient { + protocol_version: client_version.clone(), + }; + let mut server_peer_info = InitializeResult::default(); + server_peer_info.protocol_version = client_version; + + let server = serve_directly::( + ToolServer, + server_transport, + Some(client_handler.get_info()), + ); + let server_handle = tokio::spawn(async move { + server.waiting().await?; + anyhow::Ok(()) + }); + + let client = serve_directly::( + client_handler, + client_transport, + Some(server_peer_info.into()), + ); + + let result = client + .call_tool(CallToolRequestParams::new("echo")) + .await + .expect("tool call should succeed"); + + client.cancel().await.expect("client should cancel"); + server_handle.await.expect("server task").expect("server"); + result.result_type +} + +#[tokio::test] +async fn legacy_version_omits_result_type() { + assert_eq!( + call_tool_result_type(ProtocolVersion::V_2025_11_25).await, + None + ); +} + +#[tokio::test] +async fn sep_2322_version_gets_complete_result_type() { + assert_eq!( + call_tool_result_type(ProtocolVersion::V_2026_07_28).await, + Some(ResultType::COMPLETE), + ); +} diff --git a/crates/rmcp/tests/test_result_type_wire.rs b/crates/rmcp/tests/test_result_type_wire.rs new file mode 100644 index 0000000..a22933e --- /dev/null +++ b/crates/rmcp/tests/test_result_type_wire.rs @@ -0,0 +1,188 @@ +//! Wire-shape regression guards for the SEP-2322 `resultType` discriminator. +//! +//! These pin the behavior that keeps older/strict peers working: +//! - `EmptyResult` stays a bare `{}` (some peers strict-validate empty results +//! and reject extra keys), +//! - ordinary results carry `resultType: "complete"` by default, and +//! - [`ServerResult::strip_result_type_for_legacy_peer`] removes the +//! `"complete"` discriminator so legacy sessions keep their historical wire +//! shape (round-tripping a legacy result never adds the field). + +use rmcp::model::{ + CallToolResult, CompleteResult, ContentBlock, EmptyResult, GetPromptResult, ListToolsResult, + ReadResourceResult, ResultType, ServerResult, +}; +use serde_json::json; + +#[test] +fn empty_result_serializes_without_result_type() { + let value = serde_json::to_value(EmptyResult {}).expect("serialize EmptyResult"); + assert_eq!(value, json!({})); +} + +#[test] +fn call_tool_result_serializes_complete_result_type() { + let value = serde_json::to_value(CallToolResult::success(vec![ContentBlock::text("ok")])) + .expect("serialize CallToolResult"); + assert_eq!(value["resultType"], "complete"); +} + +#[test] +fn paginated_result_serializes_complete_result_type() { + let value = + serde_json::to_value(ListToolsResult::default()).expect("serialize ListToolsResult"); + assert_eq!(value["resultType"], "complete"); +} + +// Guards the convention the server handler relies on: every constructor and +// `Default` impl produces `Some(COMPLETE)`, so 2026-07-28 sessions always +// include the spec-required field unless a handler clears it explicitly. +#[test] +fn constructors_and_defaults_produce_complete_result_type() { + let complete = Some(ResultType::COMPLETE); + assert_eq!(CallToolResult::default().result_type, complete); + assert_eq!(CallToolResult::success(vec![]).result_type, complete); + assert_eq!(CallToolResult::error(vec![]).result_type, complete); + assert_eq!(CallToolResult::structured(json!({})).result_type, complete); + assert_eq!( + CallToolResult::structured_error(json!({})).result_type, + complete + ); + assert_eq!(ListToolsResult::default().result_type, complete); + assert_eq!( + ListToolsResult::with_all_items(vec![]).result_type, + complete + ); + assert_eq!(ReadResourceResult::new(vec![]).result_type, complete); + assert_eq!(GetPromptResult::default().result_type, complete); + assert_eq!(GetPromptResult::new(vec![]).result_type, complete); + assert_eq!(CompleteResult::default().result_type, complete); +} + +#[test] +fn absent_result_type_deserializes_as_none() { + let legacy = json!({ + "content": [], + "isError": false, + }); + + let result: CallToolResult = + serde_json::from_value(legacy).expect("deserialize legacy CallToolResult"); + + assert_eq!(result.result_type, None); +} + +#[test] +fn legacy_call_tool_result_round_trips_without_result_type() { + let legacy = json!({ + "content": [], + "isError": false, + }); + + let result: CallToolResult = + serde_json::from_value(legacy.clone()).expect("deserialize legacy CallToolResult"); + let reserialized = serde_json::to_value(result).expect("serialize CallToolResult"); + + assert_eq!(reserialized, legacy); +} + +#[test] +fn call_tool_result_should_reject_non_complete_result_type() { + let input_required = json!({ + "resultType": "input_required", + "_meta": { + "example.com/replica": "r1", + }, + }); + + assert!(serde_json::from_value::(input_required).is_err()); +} + +#[test] +fn server_result_should_preserve_input_required_result_when_meta_present() { + let input_required = json!({ + "resultType": "input_required", + "requestState": "sealed", + "_meta": { + "example.com/replica": "r1", + }, + }); + + let result = + serde_json::from_value::(input_required).expect("deserialize ServerResult"); + + match result { + ServerResult::InputRequiredResult(result) => { + assert_eq!( + (result.request_state.as_deref(), result.meta.is_some()), + (Some("sealed"), true) + ); + } + _ => panic!("expected InputRequiredResult"), + } +} + +#[test] +fn server_result_should_accept_legacy_call_tool_result_when_only_meta_present() { + let legacy = json!({ + "_meta": { + "example.com/replica": "r1", + }, + }); + + let result = + serde_json::from_value::(legacy).expect("deserialize legacy CallToolResult"); + + match result { + ServerResult::CallToolResult(result) => { + assert_eq!( + ( + result.result_type, + result.content.is_empty(), + result.meta.is_some() + ), + (None, true, true) + ); + } + _ => panic!("expected CallToolResult"), + } +} + +#[test] +fn strip_removes_complete_result_type() { + let mut result = + ServerResult::CallToolResult(CallToolResult::success(vec![ContentBlock::text("ok")])); + result.strip_result_type_for_legacy_peer(); + + let value = serde_json::to_value(result).expect("serialize ServerResult"); + assert_eq!( + value, + json!({ + "content": [{ "type": "text", "text": "ok" }], + "isError": false, + }) + ); +} + +#[test] +fn strip_removes_complete_result_type_from_paginated_result() { + let mut result = ServerResult::ListToolsResult(ListToolsResult::default()); + result.strip_result_type_for_legacy_peer(); + + let value = serde_json::to_value(result).expect("serialize ServerResult"); + assert_eq!(value, json!({ "tools": [] })); +} + +#[test] +fn strip_preserves_custom_result_type() { + let streaming: ResultType = + serde_json::from_value(json!("streaming")).expect("deserialize ResultType"); + let mut tool_result = CallToolResult::success(vec![ContentBlock::text("ok")]); + tool_result.result_type = Some(streaming); + + let mut result = ServerResult::CallToolResult(tool_result); + result.strip_result_type_for_legacy_peer(); + + let value = serde_json::to_value(result).expect("serialize ServerResult"); + assert_eq!(value["resultType"], "streaming"); +} diff --git a/crates/rmcp/tests/test_sampling.rs b/crates/rmcp/tests/test_sampling.rs index 83d3f6f..3c82fd4 100644 --- a/crates/rmcp/tests/test_sampling.rs +++ b/crates/rmcp/tests/test_sampling.rs @@ -9,6 +9,7 @@ use rmcp::{ model::*, service::{RequestContext, Service}, }; +use rstest::rstest; #[tokio::test] async fn test_basic_sampling_message_creation() -> Result<()> { @@ -370,6 +371,39 @@ fn test_tool_result_content_requires_content() { assert!(err.to_string().contains("missing field `content`")); } +#[rstest] +#[case::array(serde_json::json!([{ "city": "SF", "temp": 72 }, { "city": "NY", "temp": 65 }]))] +#[case::string(serde_json::json!("sunny"))] +#[case::integer(serde_json::json!(42))] +#[case::float(serde_json::json!(3.5))] +#[case::boolean(serde_json::json!(true))] +fn tool_result_content_round_trips_non_object_structured_content( + #[case] structured: serde_json::Value, +) -> Result<()> { + let mut tool_result = ToolResultContent::new("call_123", vec![ContentBlock::text("x")]); + tool_result.structured_content = Some(structured); + + let json = serde_json::to_string(&tool_result)?; + let deserialized: ToolResultContent = serde_json::from_str(&json)?; + assert_eq!(tool_result, deserialized); + + Ok(()) +} + +// `null` is a valid JSON value, but `Option` deserializes JSON `null` +// as `None`, so `Some(Value::Null)` collapses to `None` on round-trip. +#[test] +fn tool_result_content_null_structured_content_round_trips_to_none() -> Result<()> { + let mut tool_result = ToolResultContent::new("call_123", vec![ContentBlock::text("x")]); + tool_result.structured_content = Some(serde_json::Value::Null); + + let json = serde_json::to_string(&tool_result)?; + let deserialized: ToolResultContent = serde_json::from_str(&json)?; + assert_eq!(deserialized.structured_content, None); + + Ok(()) +} + #[tokio::test] async fn test_sampling_message_with_tool_use() -> Result<()> { let message = SamplingMessage::assistant_tool_use( diff --git a/crates/rmcp/tests/test_sep_2260_request_association.rs b/crates/rmcp/tests/test_sep_2260_request_association.rs new file mode 100644 index 0000000..787d634 --- /dev/null +++ b/crates/rmcp/tests/test_sep_2260_request_association.rs @@ -0,0 +1,298 @@ +#![cfg(all(feature = "server", feature = "client", not(feature = "local")))] +#![expect( + deprecated, + reason = "This test verifies request association for the deprecated sampling API" +)] + +use std::sync::{Arc, Mutex}; + +use rmcp::{ + ClientHandler, RoleClient, RoleServer, ServerHandler, ServiceError, ServiceExt, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ClientConfig, ContentBlock, + CreateMessageRequest, CreateMessageRequestParams, CreateMessageResult, ProtocolVersion, + SamplingMessage, ServerCapabilities, ServerConfig, ServerRequest, + }, + service::{RequestContext, RunningService, serve_directly}, +}; +use serde_json::{Value, json}; +use tokio::{ + io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream, Lines, ReadHalf, WriteHalf}, + sync::oneshot, +}; + +type RequestResultSender = oneshot::Sender>; + +#[derive(Clone)] +struct SamplingServer { + outside: Arc>>, +} + +impl ServerHandler for SamplingServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + let peer = context.peer.clone(); + let slot = self.outside.clone(); + + let use_generic = request.name == "sample_generic"; + tokio::spawn(async move { + let outside = if use_generic { + peer.send_request(ServerRequest::CreateMessageRequest( + CreateMessageRequest::new(CreateMessageRequestParams::new( + vec![SamplingMessage::user_text("standalone-generic")], + 16, + )), + )) + .await + .map(|_| ()) + } else { + peer.create_message(CreateMessageRequestParams::new( + vec![SamplingMessage::user_text("standalone")], + 16, + )) + .await + .map(|_| ()) + }; + if let Some(tx) = slot.lock().unwrap().take() { + let _ = tx.send(outside); + } + }); + + let nested = context + .peer + .create_message(CreateMessageRequestParams::new( + vec![SamplingMessage::user_text("nested")], + 16, + )) + .await; + nested.map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?; + Ok(CallToolResult::success(vec![ContentBlock::text("ok")]).into()) + } +} + +#[derive(Clone)] +struct SamplingClient; + +impl ClientHandler for SamplingClient { + async fn create_message( + &self, + _params: CreateMessageRequestParams, + _context: RequestContext, + ) -> Result { + Ok(CreateMessageResult::new( + SamplingMessage::assistant_text("pong"), + "test-model".to_string(), + ) + .with_stop_reason(CreateMessageResult::STOP_REASON_END_TURN)) + } + + fn get_info(&self) -> ClientConfig { + let mut info = ClientConfig::default(); + info.protocol_version = ProtocolVersion::V_2026_07_28; + info + } +} + +/// Connects the pair on `2026-07-28`. That revision dropped the `initialize` +/// handshake, so the version is agreed up front the way a discover-lifecycle +/// startup leaves it. +fn serve_modern_pair( + server: SamplingServer, +) -> ( + RunningService, + RunningService, +) { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server_peer_info = server.get_info(); + server_peer_info.protocol_version = ProtocolVersion::V_2026_07_28; + + let running_server = serve_directly::( + server, + server_transport, + Some(SamplingClient.get_info()), + ); + let client = serve_directly::( + SamplingClient, + client_transport, + Some(server_peer_info.into()), + ); + (running_server, client) +} + +#[tokio::test] +async fn nested_sampling_allowed_standalone_rejected() -> anyhow::Result<()> { + let (tx, rx) = oneshot::channel(); + let server = SamplingServer { + outside: Arc::new(Mutex::new(Some(tx))), + }; + let (running_server, client) = serve_modern_pair(server); + let server_handle = tokio::spawn(async move { + running_server.waiting().await?; + anyhow::Ok(()) + }); + + let result = client + .peer() + .call_tool(CallToolRequestParams::new("sample")) + .await?; + assert_eq!( + result.content.first().unwrap().as_text().unwrap().text, + "ok" + ); + + let outside = rx.await?; + assert!(matches!(outside, Err(ServiceError::McpError(_)))); + + client.cancel().await?; + let _ = server_handle.await?; + Ok(()) +} + +#[tokio::test] +async fn generic_send_request_bypass_rejected() -> anyhow::Result<()> { + let (tx, rx) = oneshot::channel(); + let server = SamplingServer { + outside: Arc::new(Mutex::new(Some(tx))), + }; + let (running_server, client) = serve_modern_pair(server); + let server_handle = tokio::spawn(async move { + running_server.waiting().await?; + anyhow::Ok(()) + }); + + let result = client + .peer() + .call_tool(CallToolRequestParams::new("sample_generic")) + .await?; + assert_eq!( + result.content.first().unwrap().as_text().unwrap().text, + "ok" + ); + + let outside = rx.await?; + assert!( + matches!(outside, Err(ServiceError::McpError(_))), + "generic send_request must not bypass SEP-2260 enforcement" + ); + + client.cancel().await?; + let _ = server_handle.await?; + Ok(()) +} + +// A compliant rmcp server cannot produce an unassociated server-to-client +// request at >= 2026-07-28 (send-side enforcement blocks it), so the client's +// receive-side enforcement is exercised with a raw JSON-RPC server. +type RawServer = ( + Lines>>, + WriteHalf, +); + +async fn raw_initialize(io: DuplexStream, protocol_version: &str) -> anyhow::Result { + let (read, mut write) = tokio::io::split(io); + let mut lines = BufReader::new(read).lines(); + let init: Value = serde_json::from_str(&lines.next_line().await?.expect("initialize request"))?; + assert_eq!(init["method"], "initialize"); + let response = json!({ + "jsonrpc": "2.0", + "id": init["id"], + "result": { + "protocolVersion": protocol_version, + "capabilities": {}, + "serverInfo": { "name": "raw-server", "version": "0.0.0" } + } + }); + write.write_all(format!("{response}\n").as_bytes()).await?; + let initialized: Value = + serde_json::from_str(&lines.next_line().await?.expect("initialized notification"))?; + assert_eq!(initialized["method"], "notifications/initialized"); + Ok((lines, write)) +} + +async fn raw_request(server: &mut RawServer, request: Value) -> anyhow::Result { + let (lines, write) = server; + write.write_all(format!("{request}\n").as_bytes()).await?; + Ok(serde_json::from_str( + &lines.next_line().await?.expect("response"), + )?) +} + +fn raw_sampling_request(id: u32) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id, + "method": "sampling/createMessage", + "params": { + "messages": [{ "role": "user", "content": { "type": "text", "text": "hi" } }], + "maxTokens": 16 + } + }) +} + +#[tokio::test] +async fn unassociated_server_request_rejected_with_invalid_params() -> anyhow::Result<()> { + let (client_io, server_io) = tokio::io::duplex(4096); + let raw = tokio::spawn(async move { + let mut server = raw_initialize(server_io, "2026-07-28").await?; + raw_request(&mut server, raw_sampling_request(100)).await + }); + + let client = SamplingClient.serve(client_io).await?; + let response = raw.await??; + assert_eq!( + response["error"]["code"], -32602, + "SEP-2260: unassociated server-to-client request must be rejected with invalid params, got {response}" + ); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn unassociated_server_request_allowed_on_legacy_protocol() -> anyhow::Result<()> { + let (client_io, server_io) = tokio::io::duplex(4096); + let raw = tokio::spawn(async move { + let mut server = raw_initialize(server_io, "2025-11-25").await?; + raw_request(&mut server, raw_sampling_request(100)).await + }); + + let client = SamplingClient.serve(client_io).await?; + let response = raw.await??; + assert_eq!( + response["result"]["model"], "test-model", + "pre-2026-07-28 peers keep the permissive behavior, got {response}" + ); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn unassociated_ping_allowed() -> anyhow::Result<()> { + let (client_io, server_io) = tokio::io::duplex(4096); + let raw = tokio::spawn(async move { + let mut server = raw_initialize(server_io, "2026-07-28").await?; + raw_request( + &mut server, + json!({ "jsonrpc": "2.0", "id": 101, "method": "ping" }), + ) + .await + }); + + let client = SamplingClient.serve(client_io).await?; + let response = raw.await??; + assert!( + response.get("error").is_none(), + "SEP-2260 excepts ping from request association, got {response}" + ); + + client.cancel().await?; + Ok(()) +} diff --git a/crates/rmcp/tests/test_sep_2260_stream_enforcement.rs b/crates/rmcp/tests/test_sep_2260_stream_enforcement.rs new file mode 100644 index 0000000..72a0832 --- /dev/null +++ b/crates/rmcp/tests/test_sep_2260_stream_enforcement.rs @@ -0,0 +1,309 @@ +//! SEP-2260 follow-up (#1033): stream-based receive-side enforcement. +//! +//! Scripted streamable HTTP "server": answers a legacy initialize with +//! protocol 2026-07-28 AND a session id. That combination is NOT +//! spec-compliant: the 2026-07-28 revision removes protocol-level sessions +//! and the standalone GET stream (SEP-2567; transports spec: "do not mint +//! or echo session IDs"). Receive-side enforcement (#1033) exists precisely +//! to protect the client from non-conforming servers, and rmcp's client +//! tolerates the session id and opens the standalone GET stream — so this +//! is the reachable path where the client has BOTH a GET stream and strict +//! SEP-2260 enforcement. +#![cfg(all( + feature = "client", + feature = "transport-streamable-http-client", + not(feature = "local") +))] +#![expect( + deprecated, + reason = "Sampling is deprecated by SEP-2577 but remains the canonical restricted request" +)] + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use futures::{StreamExt, stream::BoxStream}; +use http::{HeaderName, HeaderValue}; +use rmcp::{ + ClientHandler, + model::{ + ClientConfig, ClientJsonRpcMessage, CreateMessageRequestParams, CreateMessageResult, + ProtocolVersion, SamplingMessage, ServerCapabilities, ServerConfig, ServerJsonRpcMessage, + }, + service::{ClientLifecycleMode, RequestContext, RoleClient, serve_client_with_lifecycle}, + transport::streamable_http_client::{ + StreamableHttpClient, StreamableHttpClientTransport, StreamableHttpClientTransportConfig, + StreamableHttpError, StreamableHttpPostResponse, + }, +}; +use serde_json::{Value, json}; +use sse_stream::{Error as SseError, Sse}; +use tokio::sync::{Mutex, mpsc}; + +fn to_sse(message: Value) -> Result { + Ok(Sse { + event: None, + data: Some(message.to_string()), + id: None, + retry: None, + }) +} + +fn message_stream(rx: mpsc::Receiver) -> BoxStream<'static, Result> { + tokio_stream::wrappers::ReceiverStream::new(rx) + .map(to_sse) + .boxed() +} + +/// Scripted server: initialize -> JSON init result (2026-07-28 + session); +/// first non-initialize request POST -> SSE stream fed by `post_stream`; +/// everything else -> Accepted. Every message the client POSTs is forwarded +/// to `posted`. +#[derive(Clone)] +struct ScriptedServer { + get_stream: Arc>>>, + post_stream: Arc>>>, + posted: mpsc::UnboundedSender, +} + +impl StreamableHttpClient for ScriptedServer { + type Error = std::io::Error; + + async fn post_message( + &self, + _uri: Arc, + message: ClientJsonRpcMessage, + _session_id: Option>, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + let value = serde_json::to_value(&message).expect("serialize client message"); + // Receiver drop is normal at test teardown; never panic in the transport task. + let _ = self.posted.send(value.clone()); + if value["method"] == "initialize" { + let mut info = ServerConfig::new(ServerCapabilities::default()); + info.protocol_version = ProtocolVersion::V_2026_07_28; + let response = ServerJsonRpcMessage::response( + rmcp::model::ServerResult::InitializeResult(info), + serde_json::from_value(value["id"].clone()).expect("request id"), + ); + return Ok(StreamableHttpPostResponse::Json( + response, + Some("scripted-session".into()), + )); + } + if matches!(message, ClientJsonRpcMessage::Request(_)) { + // Fail as a transport error rather than panicking: this code runs + // in the transport task, where a panic is swallowed and shows up + // only as an opaque timeout in the test. + let rx = self.post_stream.lock().await.take().ok_or_else(|| { + StreamableHttpError::Client(std::io::Error::other( + "scripted server expects exactly one non-initialize request POST", + )) + })?; + return Ok(StreamableHttpPostResponse::Sse(message_stream(rx), None)); + } + Ok(StreamableHttpPostResponse::Accepted) + } + + async fn delete_session( + &self, + _uri: Arc, + _session_id: Arc, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + Ok(()) + } + + async fn get_stream( + &self, + _uri: Arc, + _session_id: Option>, + _last_event_id: Option, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result>, StreamableHttpError> { + match self.get_stream.lock().await.take() { + Some(rx) => Ok(message_stream(rx)), + // Reconnect after the scripted stream ends: stay silent. + None => Ok(futures::stream::pending().boxed()), + } + } +} + +#[derive(Clone)] +struct SamplingClient { + // Every invocation is recorded here. Tests assert on the receiver (the + // negative test asserts it stays empty); a panic in this handler would + // run in a spawned task and be silently swallowed. + sampled: mpsc::UnboundedSender, +} + +impl ClientHandler for SamplingClient { + async fn create_message( + &self, + params: CreateMessageRequestParams, + _context: RequestContext, + ) -> Result { + let _ = self.sampled.send(params); + Ok(CreateMessageResult::new( + SamplingMessage::assistant_text("pong"), + "test-model".to_string(), + )) + } + + fn get_info(&self) -> ClientConfig { + ClientConfig::default() + } +} + +fn sampling_request(id: u32) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id, + "method": "sampling/createMessage", + "params": { + "messages": [{ "role": "user", "content": { "type": "text", "text": "hi" } }], + "maxTokens": 16 + } + }) +} + +fn tools_list_response(id: &Value) -> Value { + json!({ "jsonrpc": "2.0", "id": id, "result": { "tools": [] } }) +} + +async fn next_posted(posted: &mut mpsc::UnboundedReceiver) -> Value { + tokio::time::timeout(Duration::from_secs(5), posted.recv()) + .await + .expect("posted message within 5s") + .expect("channel open") +} + +struct Harness { + client: rmcp::service::RunningService, + /// Every message the client POSTs to the scripted server. + posted: mpsc::UnboundedReceiver, + /// Feeds the standalone GET stream. + get_tx: mpsc::Sender, + /// Feeds the SSE stream of the in-flight tools/list POST. + post_tx: mpsc::Sender, + /// In-flight tools/list call (response withheld until the test releases it). + call: tokio::task::JoinHandle>, + tools_list_id: Value, + /// Sampling params seen by the client handler. + sampled: mpsc::UnboundedReceiver, +} + +/// Drive startup + one in-flight tools/list. +async fn setup() -> Harness { + let (get_tx, get_rx) = mpsc::channel(8); + let (post_tx, post_rx) = mpsc::channel(8); + let (posted_tx, mut posted_rx) = mpsc::unbounded_channel(); + let (sampled_tx, sampled_rx) = mpsc::unbounded_channel(); + let server = ScriptedServer { + get_stream: Arc::new(Mutex::new(Some(get_rx))), + post_stream: Arc::new(Mutex::new(Some(post_rx))), + posted: posted_tx, + }; + let transport = StreamableHttpClientTransport::with_client( + server, + StreamableHttpClientTransportConfig::with_uri("http://scripted/mcp"), + ); + let client = serve_client_with_lifecycle( + SamplingClient { + sampled: sampled_tx, + }, + transport, + ClientLifecycleMode::Initialize, + ) + .await + .expect("initialize against scripted server"); + + // initialize + notifications/initialized already posted during startup. + assert_eq!(next_posted(&mut posted_rx).await["method"], "initialize"); + assert_eq!( + next_posted(&mut posted_rx).await["method"], + "notifications/initialized" + ); + + // Unrelated outbound request, kept in flight (response withheld). + let peer = client.peer().clone(); + let call = tokio::spawn(async move { peer.list_tools(None).await }); + let tools_list = next_posted(&mut posted_rx).await; + assert_eq!(tools_list["method"], "tools/list"); + let tools_list_id = tools_list["id"].clone(); + + Harness { + client, + posted: posted_rx, + get_tx, + post_tx, + call, + tools_list_id, + sampled: sampled_rx, + } +} + +/// #1033 scenario 1: a restricted request on the standalone GET stream while +/// an unrelated outbound request is in flight must be rejected with -32602. +/// (The coarse check from #1029 incorrectly accepted this.) +#[tokio::test] +async fn restricted_request_on_get_stream_rejected_while_unrelated_request_in_flight() +-> anyhow::Result<()> { + let mut h = setup().await; + + h.get_tx.send(sampling_request(100)).await?; + + let rejection = next_posted(&mut h.posted).await; + assert_eq!( + rejection["id"], 100, + "reply to the sampling request: {rejection}" + ); + assert_eq!( + rejection["error"]["code"], -32602, + "SEP-2260: GET-stream request must be rejected even with an unrelated \ + request in flight, got {rejection}" + ); + + h.post_tx + .send(tools_list_response(&h.tools_list_id)) + .await?; + tokio::time::timeout(Duration::from_secs(5), h.call).await???; + // Rejected means rejected: the handler must not ALSO have been invoked. + assert!( + h.sampled.try_recv().is_err(), + "handler must not see the rejected sampling request" + ); + tokio::time::timeout(Duration::from_secs(5), h.client.cancel()).await??; + Ok(()) +} + +/// Positive twin: the same restricted request arriving on the SSE stream of +/// the originating POST is dispatched to the handler and answered. +#[tokio::test] +async fn restricted_request_on_originating_post_stream_is_dispatched() -> anyhow::Result<()> { + let mut h = setup().await; + + h.post_tx.send(sampling_request(200)).await?; + + let response = next_posted(&mut h.posted).await; + assert_eq!( + response["id"], 200, + "reply to the sampling request: {response}" + ); + assert_eq!( + response["result"]["model"], "test-model", + "request on the originating POST stream must reach the handler, got {response}" + ); + tokio::time::timeout(Duration::from_secs(5), h.sampled.recv()) + .await? + .expect("handler invoked"); + + h.post_tx + .send(tools_list_response(&h.tools_list_id)) + .await?; + tokio::time::timeout(Duration::from_secs(5), h.call).await???; + tokio::time::timeout(Duration::from_secs(5), h.client.cancel()).await??; + Ok(()) +} diff --git a/crates/rmcp/tests/test_sep_2260_stream_routing.rs b/crates/rmcp/tests/test_sep_2260_stream_routing.rs new file mode 100644 index 0000000..e883ccd --- /dev/null +++ b/crates/rmcp/tests/test_sep_2260_stream_routing.rs @@ -0,0 +1,165 @@ +//! SEP-2260 end-to-end: in-handler server→client requests ride the originating +//! POST's SSE stream, never the standalone GET stream. +#![cfg(not(feature = "local"))] + +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use futures::StreamExt; +use rmcp::{ + ErrorData as McpError, RoleServer, ServerHandler, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ElicitRequestParams, + ElicitationSchema, ServerCapabilities, ServerConfig, + }, + service::RequestContext, + transport::streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, + }, +}; +use serde_json::json; +use tokio_util::sync::CancellationToken; + +#[derive(Clone)] +struct ElicitingServer; + +impl ServerHandler for ElicitingServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn call_tool( + &self, + _request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + // Never answered: the test only checks the request is emitted on the right stream. + let _ = context + .peer + .create_elicitation(ElicitRequestParams::FormElicitationParams { + meta: None, + message: "need input".to_string(), + requested_schema: ElicitationSchema::new(BTreeMap::new()), + }) + .await; + Ok(CallToolResult::success(vec![ContentBlock::text("done")]).into()) + } +} + +async fn start_server(ct: CancellationToken) -> String { + let service = StreamableHttpService::new( + move || Ok(ElicitingServer), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default().with_cancellation_token(ct.child_token()), + ); + let router = axum::Router::new().nest_service("/mcp", service); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!( + "http://127.0.0.1:{}/mcp", + listener.local_addr().unwrap().port() + ); + let ct = ct.clone(); + tokio::spawn(async move { + axum::serve(listener, router) + .with_graceful_shutdown(async move { ct.cancelled().await }) + .await + .unwrap(); + }); + tokio::time::sleep(Duration::from_millis(100)).await; + url +} + +/// Read an SSE byte stream until `needle` appears or timeout. +async fn sse_contains(resp: reqwest::Response, needle: &str, timeout: Duration) -> bool { + let mut stream = resp.bytes_stream(); + tokio::time::timeout(timeout, async { + let mut buffer = String::new(); + while let Some(Ok(chunk)) = stream.next().await { + buffer.push_str(&String::from_utf8_lossy(&chunk)); + if buffer.contains(needle) { + return true; + } + } + false + }) + .await + .unwrap_or(false) +} + +#[tokio::test] +async fn elicitation_rides_originating_post_stream_not_standalone_get() { + let ct = CancellationToken::new(); + let url = start_server(ct.clone()).await; + let client = reqwest::Client::new(); + + // 2025-11-25 is the latest session-carrying version; SEP-2567 serves + // 2026-07-28+ statelessly, with no standalone GET stream to test against. + let resp = client + .post(&url) + .header("Accept", "text/event-stream, application/json") + .header("Content-Type", "application/json") + .json(&json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": { + "protocolVersion": "2025-11-25", + "capabilities": { "elicitation": {} }, + "clientInfo": { "name": "test-client", "version": "1.0.0" } + } + })) + .send() + .await + .unwrap(); + assert!(resp.status().is_success()); + let session_id = resp + .headers() + .get("Mcp-Session-Id") + .expect("session id") + .to_str() + .unwrap() + .to_string(); + + client + .post(&url) + .header("Accept", "text/event-stream, application/json") + .header("Content-Type", "application/json") + .header("Mcp-Session-Id", &session_id) + .json(&json!({"jsonrpc": "2.0", "method": "notifications/initialized"})) + .send() + .await + .unwrap(); + + // Standalone GET stream — must NEVER carry the elicitation request. + let get_stream = client + .get(&url) + .header("Accept", "text/event-stream") + .header("Mcp-Session-Id", &session_id) + .send() + .await + .unwrap(); + assert_eq!(get_stream.status(), 200); + + // The in-handler elicitation must appear on this tools/call SSE stream. + let post_stream = client + .post(&url) + .header("Accept", "text/event-stream, application/json") + .header("Content-Type", "application/json") + .header("Mcp-Session-Id", &session_id) + .json(&json!({ + "jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": { "name": "ask", "arguments": {} } + })) + .send() + .await + .unwrap(); + assert!(post_stream.status().is_success()); + + assert!( + sse_contains(post_stream, "elicitation/create", Duration::from_secs(5)).await, + "elicitation request must be delivered on the originating POST SSE stream" + ); + assert!( + !sse_contains(get_stream, "elicitation/create", Duration::from_millis(500)).await, + "standalone GET stream must not carry the elicitation request" + ); + + ct.cancel(); +} diff --git a/crates/rmcp/tests/test_server_discover.rs b/crates/rmcp/tests/test_server_discover.rs new file mode 100644 index 0000000..f2ca3df --- /dev/null +++ b/crates/rmcp/tests/test_server_discover.rs @@ -0,0 +1,255 @@ +use rmcp::model::{ + ClientCapabilities, ClientJsonRpcMessage, ClientRequest, DiscoverResult, ErrorCode, ErrorData, + JsonRpcRequest, JsonRpcResponse, ProtocolVersion, ServerJsonRpcMessage, ServerResult, +}; +use serde_json::json; + +#[test] +fn discover_request_deserializes_with_request_meta() { + let message: ClientJsonRpcMessage = serde_json::from_value(json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "server/discover", + "params": { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": { + "name": "test-client", + "version": "1.0.0" + }, + "io.modelcontextprotocol/clientCapabilities": {} + } + } + })) + .expect("discover request should deserialize"); + + let ClientJsonRpcMessage::Request(JsonRpcRequest { request, .. }) = message else { + panic!("expected request"); + }; + let ClientRequest::DiscoverRequest(request) = request else { + panic!("expected discover request"); + }; + + assert_eq!( + request + .extensions + .get::() + .and_then(|meta| meta.protocol_version()), + Some(ProtocolVersion::V_2026_07_28) + ); +} + +#[test] +fn discover_result_deserializes_to_typed_variant() { + let message: ServerJsonRpcMessage = serde_json::from_value(json!({ + "jsonrpc": "2.0", + "id": 1, + "result": { + "resultType": "complete", + "supportedVersions": ["2025-11-25", "2026-07-28"], + "capabilities": { "tools": {} }, + "serverInfo": { + "name": "test-server", + "version": "1.0.0" + }, + "ttlMs": 0, + "cacheScope": "private" + } + })) + .expect("discover result should deserialize"); + + let ServerJsonRpcMessage::Response(JsonRpcResponse { result, .. }) = message else { + panic!("expected response"); + }; + let ServerResult::DiscoverResult(DiscoverResult { + supported_versions, .. + }) = result + else { + panic!("expected discover result"); + }; + + assert_eq!( + supported_versions, + vec![ProtocolVersion::V_2025_11_25, ProtocolVersion::V_2026_07_28] + ); +} + +#[test] +fn discover_result_accepts_server_info_in_namespaced_metadata() { + let message: ServerJsonRpcMessage = serde_json::from_value(json!({ + "jsonrpc": "2.0", + "id": 1, + "result": { + "resultType": "complete", + "supportedVersions": ["2026-07-28"], + "capabilities": {}, + "ttlMs": 0, + "cacheScope": "private", + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "conformance-mock-server", + "version": "1.0.0" + }, + "unrelated": { "preserved": true } + } + } + })) + .expect("discovery response with namespaced server info should deserialize"); + + let ServerJsonRpcMessage::Response(JsonRpcResponse { result, .. }) = message else { + panic!("expected response"); + }; + let ServerResult::DiscoverResult(result) = result else { + panic!("expected discovery response, not a tool-call result"); + }; + + let server_info = result.server_info().expect("server info should be present"); + assert_eq!(server_info.name, "conformance-mock-server"); + assert_eq!(server_info.version, "1.0.0"); + + let metadata = result.meta.expect("discovery metadata should be preserved"); + assert_eq!( + metadata.0.get("io.modelcontextprotocol/serverInfo"), + Some(&json!({ + "name": "conformance-mock-server", + "version": "1.0.0" + })) + ); + assert_eq!( + metadata.0.get("unrelated"), + Some(&json!({ "preserved": true })) + ); +} + +#[test] +fn discover_result_serializes_server_info_in_namespaced_metadata() { + let mut result = DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + rmcp::model::ServerCapabilities::default(), + ); + result.meta = Some(rmcp::model::MetaObject( + json!({ + "io.modelcontextprotocol/serverInfo": { + "name": "stale-server", + "version": "0.1.0" + }, + "unrelated": { "preserved": true } + }) + .as_object() + .expect("metadata is an object") + .clone(), + )); + result.set_server_info(rmcp::model::Implementation::new("test-server", "1.0.0")); + + let serialized = serde_json::to_value(result).expect("serialize discovery result"); + assert!(serialized.get("serverInfo").is_none()); + assert_eq!( + serialized["_meta"]["io.modelcontextprotocol/serverInfo"], + json!({ + "name": "test-server", + "version": "1.0.0" + }) + ); + assert_eq!( + serialized["_meta"]["unrelated"], + json!({ "preserved": true }) + ); +} + +#[test] +fn discover_result_ignores_legacy_top_level_server_info() { + let result: DiscoverResult = serde_json::from_value(json!({ + "resultType": "complete", + "supportedVersions": ["2026-07-28"], + "capabilities": {}, + "serverInfo": { + "name": "top-level-server", + "version": "2.0.0" + }, + "ttlMs": 0, + "cacheScope": "private", + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "metadata-server", + "version": "1.0.0" + }, + "unrelated": true + } + })) + .expect("top-level server info should remain supported"); + + let server_info = result + .server_info() + .expect("namespaced server info is present"); + assert_eq!(server_info.name, "metadata-server"); + assert_eq!(server_info.version, "1.0.0"); + assert_eq!( + result + .meta + .as_ref() + .and_then(|metadata| metadata.0.get("unrelated")), + Some(&json!(true)) + ); +} + +#[test] +fn discover_result_allows_missing_or_malformed_optional_server_info() { + let result = json!({ + "resultType": "complete", + "supportedVersions": ["2026-07-28"], + "capabilities": {}, + "ttlMs": 0, + "cacheScope": "private", + "_meta": { "unrelated": true } + }); + + let result = + serde_json::from_value::(result).expect("server info metadata is optional"); + assert_eq!(result.server_info(), None); + + let malformed_server_info = json!({ + "resultType": "complete", + "supportedVersions": ["2026-07-28"], + "capabilities": {}, + "ttlMs": 0, + "cacheScope": "private", + "_meta": { "io.modelcontextprotocol/serverInfo": { "name": "missing-version" } } + }); + + let result = serde_json::from_value::(malformed_server_info) + .expect("opaque metadata should not prevent deserialization"); + assert_eq!(result.server_info(), None); +} + +#[test] +fn unsupported_protocol_version_error_matches_draft_schema() { + let error = ErrorData::unsupported_protocol_version( + ProtocolVersion::V_2026_07_28, + &[ProtocolVersion::V_2025_11_25], + ); + + assert_eq!(error.code, ErrorCode::UNSUPPORTED_PROTOCOL_VERSION); + assert_eq!( + error.data, + Some(json!({ + "requested": "2026-07-28", + "supported": ["2025-11-25"] + })) + ); +} + +#[test] +fn missing_required_capability_error_matches_draft_schema() { + let required = ClientCapabilities::builder().enable_elicitation().build(); + let error = ErrorData::missing_required_client_capability(required); + + assert_eq!(error.code, ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY); + assert_eq!( + error.data, + Some(json!({ + "requiredCapabilities": { + "elicitation": {} + } + })) + ); +} diff --git a/crates/rmcp/tests/test_server_discover_client.rs b/crates/rmcp/tests/test_server_discover_client.rs new file mode 100644 index 0000000..9046f6b --- /dev/null +++ b/crates/rmcp/tests/test_server_discover_client.rs @@ -0,0 +1,77 @@ +#![cfg(all(feature = "client", not(feature = "local")))] + +use rmcp::{ + ClientHandler, ServerHandler, ServiceExt, + model::{ + ClientCapabilities, Implementation, ProtocolVersion, RequestMetaObject, ServerCapabilities, + ServerConfig, + }, + select_protocol_version, +}; + +#[derive(Clone, Default)] +struct DiscoveryServer; + +impl ServerHandler for DiscoveryServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::new("discovery-server", "1.0.0")) + } +} + +#[derive(Clone, Default)] +struct DiscoveryClient; + +impl ClientHandler for DiscoveryClient {} + +#[test] +fn select_protocol_version_uses_client_preference_order() { + let selected = select_protocol_version( + &[ProtocolVersion::V_2026_07_28, ProtocolVersion::V_2025_11_25], + &[ProtocolVersion::V_2025_11_25, ProtocolVersion::V_2026_07_28], + ); + + assert_eq!(selected, Some(ProtocolVersion::V_2026_07_28)); +} + +#[test] +fn select_protocol_version_returns_none_without_overlap() { + let selected = select_protocol_version( + &[ProtocolVersion::V_2026_07_28], + &[ProtocolVersion::V_2025_11_25], + ); + + assert_eq!(selected, None); +} + +#[tokio::test] +async fn client_discover_helper_returns_typed_result() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + tokio::spawn(async move { + let _ = DiscoveryServer + .serve(server_transport) + .await + .expect("server should start") + .waiting() + .await; + }); + let client = DiscoveryClient + .serve(client_transport) + .await + .expect("client should connect"); + let mut meta = RequestMetaObject::new(); + meta.set_protocol_version(ProtocolVersion::V_2026_07_28); + meta.set_client_info(Implementation::new("discovery-client", "1.0.0")); + meta.set_client_capabilities(ClientCapabilities::default()); + + let result = client + .discover(meta) + .await + .expect("discover should succeed"); + + assert_eq!( + result.server_info(), + Some(Implementation::new("discovery-server", "1.0.0")) + ); + client.cancel().await.expect("client should cancel"); +} diff --git a/crates/rmcp/tests/test_server_discover_http.rs b/crates/rmcp/tests/test_server_discover_http.rs new file mode 100644 index 0000000..7762c97 --- /dev/null +++ b/crates/rmcp/tests/test_server_discover_http.rs @@ -0,0 +1,380 @@ +#![cfg(all( + not(feature = "local"), + feature = "reqwest", + feature = "transport-streamable-http-server" +))] + +use std::borrow::Cow; + +use rmcp::{ + ServerHandler, + model::{Implementation, ProtocolVersion, ServerCapabilities, ServerConfig}, + transport::streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, + }, +}; +use serde_json::json; +use tokio_util::sync::CancellationToken; + +#[derive(Clone, Default)] +struct DiscoveryServer; + +impl ServerHandler for DiscoveryServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::new("discovery-server", "1.0.0")) + .with_instructions("Use the tools carefully") + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(&[ProtocolVersion::V_2025_11_25]) + } +} + +async fn spawn_server(json_response: bool) -> (reqwest::Client, String, CancellationToken) { + spawn_server_with_legacy_session_mode(json_response, false).await +} + +async fn spawn_server_with_legacy_session_mode( + json_response: bool, + legacy_session_mode: bool, +) -> (reqwest::Client, String, CancellationToken) { + let cancellation_token = CancellationToken::new(); + let config = StreamableHttpServerConfig::default() + .with_legacy_session_mode(legacy_session_mode) + .with_json_response(json_response) + .with_sse_keep_alive(None) + .with_cancellation_token(cancellation_token.clone()); + let service: StreamableHttpService = + StreamableHttpService::new(|| Ok(DiscoveryServer), Default::default(), config); + let router = axum::Router::new().nest_service("/mcp", service); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let address = listener.local_addr().expect("listener should have address"); + + tokio::spawn({ + let cancellation_token = cancellation_token.clone(); + async move { + let _ = axum::serve(listener, router) + .with_graceful_shutdown(async move { + cancellation_token.cancelled_owned().await; + }) + .await; + } + }); + + ( + reqwest::Client::new(), + format!("http://{address}/mcp"), + cancellation_token, + ) +} + +fn discover_body(version: Option<&str>) -> serde_json::Value { + let meta = version.map(|version| { + json!({ + "io.modelcontextprotocol/protocolVersion": version, + "io.modelcontextprotocol/clientInfo": { + "name": "test-client", + "version": "1.0.0" + }, + "io.modelcontextprotocol/clientCapabilities": {} + }) + }); + json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "server/discover", + "params": meta.map(|meta| json!({ "_meta": meta })).unwrap_or_else(|| json!({})) + }) +} + +async fn post_discover( + client: &reqwest::Client, + url: &str, + header_version: &str, + body_version: Option<&str>, +) -> reqwest::Response { + client + .post(url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", header_version) + .header("Mcp-Method", "server/discover") + .json(&discover_body(body_version)) + .send() + .await + .expect("discover request should send") +} + +#[tokio::test] +async fn discover_returns_server_metadata_without_session() { + let (client, url, cancellation_token) = spawn_server(true).await; + + let response = post_discover(&client, &url, "2025-11-25", Some("2025-11-25")).await; + + assert_eq!(response.status(), 200); + let body: serde_json::Value = response.json().await.expect("response should be JSON"); + assert_eq!( + body["result"], + json!({ + "resultType": "complete", + "supportedVersions": ["2025-11-25"], + "capabilities": { "tools": {} }, + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "discovery-server", + "version": "1.0.0" + } + }, + "instructions": "Use the tools carefully", + "ttlMs": 0, + "cacheScope": "private" + }) + ); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn discover_does_not_require_initialization_in_legacy_session_mode() { + let (client, url, cancellation_token) = spawn_server_with_legacy_session_mode(true, true).await; + + let response = post_discover(&client, &url, "2025-11-25", Some("2025-11-25")).await; + + assert_eq!(response.status(), 200); + assert!(response.headers().get("Mcp-Session-Id").is_none()); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn discover_rejects_unsupported_version_with_http_400() { + let (client, url, cancellation_token) = spawn_server(true).await; + + let response = post_discover(&client, &url, "2026-07-28", Some("2026-07-28")).await; + + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await.expect("response should be JSON"); + assert_eq!(body["error"]["code"], -32022); + assert_eq!( + body["error"]["data"], + json!({ + "requested": "2026-07-28", + "supported": ["2025-11-25"] + }) + ); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn discover_rejects_unknown_version_with_typed_error() { + let (client, url, cancellation_token) = spawn_server(true).await; + + let response = post_discover(&client, &url, "2099-01-01", Some("2099-01-01")).await; + + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await.expect("response should be JSON"); + assert_eq!(body["error"]["code"], -32022); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn regular_request_rejects_server_unsupported_meta_version() { + let (client, url, cancellation_token) = spawn_server(true).await; + let body = json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28" + } + } + }); + + let response = client + .post(&url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2026-07-28") + .header("Mcp-Method", "tools/list") + .json(&body) + .send() + .await + .expect("request should send"); + + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await.expect("response should be JSON"); + assert_eq!(body["error"]["code"], -32022); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn unknown_rpc_uses_http_404_for_per_request_protocol() { + let (client, url, cancellation_token) = spawn_server(true).await; + let body = json!({ + "jsonrpc": "2.0", + "id": 3, + "method": "unknown/method", + "params": { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2025-11-25" + } + } + }); + + let response = client + .post(&url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2025-11-25") + .json(&body) + .send() + .await + .expect("request should send"); + + assert_eq!(response.status(), 404); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn legacy_unknown_rpc_preserves_http_200_jsonrpc_error() { + let (client, url, cancellation_token) = spawn_server(true).await; + let body = json!({ + "jsonrpc": "2.0", + "id": 5, + "method": "unknown/method", + "params": {} + }); + + let response = client + .post(&url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2025-11-25") + .json(&body) + .send() + .await + .expect("request should send"); + + assert_eq!(response.status(), 200); + let body: serde_json::Value = response.json().await.expect("response should be JSON"); + assert_eq!(body["error"]["code"], -32601); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn discover_rejects_header_meta_version_mismatch() { + let (client, url, cancellation_token) = spawn_server(true).await; + + let response = post_discover(&client, &url, "2026-07-28", Some("2025-11-25")).await; + + assert_eq!(response.status(), 400); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn discover_rejects_missing_request_meta() { + let (client, url, cancellation_token) = spawn_server(true).await; + + let response = post_discover(&client, &url, "2025-11-25", None).await; + + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await.expect("response should be JSON"); + assert_eq!(body["error"]["code"], -32602); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn discover_rejects_missing_client_capabilities() { + let (client, url, cancellation_token) = spawn_server(true).await; + let body = json!({ + "jsonrpc": "2.0", + "id": 4, + "method": "server/discover", + "params": { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2025-11-25", + "io.modelcontextprotocol/clientInfo": { + "name": "test-client", + "version": "1.0.0" + } + } + } + }); + + let response = client + .post(&url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2025-11-25") + .json(&body) + .send() + .await + .expect("request should send"); + + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await.expect("response should be JSON"); + assert_eq!(body["error"]["code"], -32602); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn discover_accepts_missing_optional_client_info() { + let (client, url, cancellation_token) = spawn_server(true).await; + let body = json!({ + "jsonrpc": "2.0", + "id": 5, + "method": "server/discover", + "params": { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2025-11-25", + "io.modelcontextprotocol/clientCapabilities": {} + } + } + }); + + let response = client + .post(&url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2025-11-25") + .json(&body) + .send() + .await + .expect("request should send"); + + assert_eq!(response.status(), 200); + let body: serde_json::Value = response.json().await.expect("response should be JSON"); + assert!(body.get("result").is_some()); + + cancellation_token.cancel(); +} + +#[tokio::test] +async fn discover_error_uses_http_400_when_sse_is_configured() { + let (client, url, cancellation_token) = spawn_server(false).await; + + let response = post_discover(&client, &url, "2026-07-28", Some("2026-07-28")).await; + + assert_eq!(response.status(), 400); + assert_eq!( + response + .headers() + .get("Content-Type") + .and_then(|value| value.to_str().ok()), + Some("application/json") + ); + + cancellation_token.cancel(); +} diff --git a/crates/rmcp/tests/test_server_initialization.rs b/crates/rmcp/tests/test_server_initialization.rs index 6542e29..5be178e 100644 --- a/crates/rmcp/tests/test_server_initialization.rs +++ b/crates/rmcp/tests/test_server_initialization.rs @@ -6,7 +6,7 @@ use common::handlers::TestServer; use rmcp::{ ServerHandler, ServiceExt, model::{ - ClientJsonRpcMessage, ProtocolVersion, ServerCapabilities, ServerInfo, + ClientJsonRpcMessage, ProtocolVersion, ServerCapabilities, ServerConfig, ServerJsonRpcMessage, ServerResult, }, transport::{IntoTransport, Transport}, @@ -281,8 +281,8 @@ async fn server_falls_back_when_client_protocol_version_unknown() { struct PinnedServer; impl ServerHandler for PinnedServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().build()) .with_protocol_version(ProtocolVersion::V_2025_06_18) } } diff --git a/crates/rmcp/tests/test_sse_concurrent_streams.rs b/crates/rmcp/tests/test_sse_concurrent_streams.rs index e1e8852..23ab8d0 100644 --- a/crates/rmcp/tests/test_sse_concurrent_streams.rs +++ b/crates/rmcp/tests/test_sse_concurrent_streams.rs @@ -18,7 +18,7 @@ use std::time::Duration; use futures::StreamExt; use rmcp::{ RoleServer, ServerHandler, - model::{Implementation, ServerCapabilities, ServerInfo, ToolsCapability}, + model::{Implementation, ServerCapabilities, ServerConfig, ToolsCapability}, service::NotificationContext, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, @@ -45,8 +45,8 @@ impl TestServer { } impl ServerHandler for TestServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools_with({ let mut tools = ToolsCapability::default(); diff --git a/crates/rmcp/tests/test_stateless_protocol_version.rs b/crates/rmcp/tests/test_stateless_protocol_version.rs index 5103ddd..cc2a4b5 100644 --- a/crates/rmcp/tests/test_stateless_protocol_version.rs +++ b/crates/rmcp/tests/test_stateless_protocol_version.rs @@ -1,33 +1,99 @@ //! Tests for protocol version negotiation in stateless HTTP mode. //! -//! Known versions are echoed back; unknown versions fall back to LATEST. +//! Supported handshake versions are echoed back; unknown versions, versions +//! outside the server's `supported_protocol_versions`, and versions that no +//! longer have an `initialize` handshake fall back to the handler's own +//! version. #![cfg(not(feature = "local"))] +use std::borrow::Cow; + use rmcp::{ - model::ProtocolVersion, + ErrorData, RoleServer, ServerHandler, + model::{ + InitializeRequestParams, InitializeResult, ProtocolVersion, ServerCapabilities, + ServerConfig, + }, + service::RequestContext, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, }; use tokio_util::sync::CancellationToken; -mod common; -use common::calculator::Calculator; +#[derive(Clone, Default)] +struct OverridingInitialize; -fn stateless_json_config() -> StreamableHttpServerConfig { +impl ServerHandler for OverridingInitialize { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::default()) + } + + async fn initialize( + &self, + _request: InitializeRequestParams, + _context: RequestContext, + ) -> Result { + Ok(self.get_info()) + } +} + +/// Every known version whose lifecycle still runs the `initialize` handshake. +/// `2026-07-28` replaced the handshake with per-request metadata, so this is +/// also the list a server that has not implemented that revision supports. +const HANDSHAKE_VERSIONS: &[ProtocolVersion] = &[ + ProtocolVersion::V_2024_11_05, + ProtocolVersion::V_2025_03_26, + ProtocolVersion::V_2025_06_18, + ProtocolVersion::V_2025_11_25, +]; + +/// Overrides `initialize`, so the handler-side default negotiation never runs, +/// *and* narrows the supported versions. +#[derive(Clone, Default)] +struct NarrowedOverridingInitialize; + +impl ServerHandler for NarrowedOverridingInitialize { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::default()) + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(HANDSHAKE_VERSIONS) + } + + async fn initialize( + &self, + _request: InitializeRequestParams, + _context: RequestContext, + ) -> Result { + Ok(self.get_info()) + } +} + +fn stateless_sse_config() -> StreamableHttpServerConfig { StreamableHttpServerConfig::default() - .with_stateful_mode(false) - .with_json_response(true) + .with_legacy_session_mode(false) .with_sse_keep_alive(None) .with_cancellation_token(CancellationToken::new()) } +fn stateless_json_config() -> StreamableHttpServerConfig { + stateless_sse_config().with_json_response(true) +} + async fn spawn_server( config: StreamableHttpServerConfig, +) -> (reqwest::Client, String, CancellationToken) { + spawn_server_of::(config).await +} + +async fn spawn_server_of( + config: StreamableHttpServerConfig, ) -> (reqwest::Client, String, CancellationToken) { let ct = config.cancellation_token.clone(); - let service: StreamableHttpService = - StreamableHttpService::new(|| Ok(Calculator::new()), Default::default(), config); + let service: StreamableHttpService = + StreamableHttpService::new(|| Ok(H::default()), Default::default(), config); let router = axum::Router::new().nest_service("/mcp", service); let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -65,19 +131,35 @@ async fn post_init(client: &reqwest::Client, url: &str, body_version: &str) -> s .await .expect("send request"); assert!(resp.status().is_success(), "HTTP {}", resp.status()); - resp.json().await.expect("parse JSON") + let is_json = resp + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.starts_with("application/json")); + if is_json { + resp.json().await.expect("parse JSON") + } else { + let body = resp.text().await.expect("read SSE body"); + let data = body + .lines() + .find_map(|line| line.strip_prefix("data:")) + .map(str::trim) + .filter(|data| !data.is_empty()) + .expect("SSE response contains data"); + serde_json::from_str(data).expect("parse SSE data") + } } #[tokio::test] -async fn stateless_init_echoes_known_version() { +async fn stateless_json_init_echoes_handshake_versions_when_handler_overrides_initialize() { let (client, url, ct) = spawn_server(stateless_json_config()).await; - for version in ProtocolVersion::KNOWN_VERSIONS { + for version in HANDSHAKE_VERSIONS { let resp = post_init(&client, &url, version.as_str()).await; assert_eq!( resp["result"]["protocolVersion"], version.as_str(), - "known version {version} should be echoed back" + "handshake version {version} should be echoed back" ); } @@ -85,14 +167,62 @@ async fn stateless_init_echoes_known_version() { } #[tokio::test] -async fn stateless_init_unknown_version_falls_back_to_latest() { +async fn stateless_sse_init_echoes_handshake_versions_when_handler_overrides_initialize() { + let (client, url, ct) = spawn_server(stateless_sse_config()).await; + + for version in HANDSHAKE_VERSIONS { + let resp = post_init(&client, &url, version.as_str()).await; + assert_eq!( + resp["result"]["protocolVersion"], + version.as_str(), + "handshake version {version} should be echoed back" + ); + } + + ct.cancel(); +} + +#[tokio::test] +async fn stateless_json_init_preserves_handler_fallback_for_unknown_version() { let (client, url, ct) = spawn_server(stateless_json_config()).await; let resp = post_init(&client, &url, "1999-01-01").await; assert_eq!( resp["result"]["protocolVersion"], ProtocolVersion::LATEST.as_str(), - "unknown version should fall back to LATEST" + "unknown version should preserve the handler's fallback" + ); + + ct.cancel(); +} + +#[tokio::test] +async fn stateless_json_init_echoes_versions_the_server_narrowed_to() { + let (client, url, ct) = + spawn_server_of::(stateless_json_config()).await; + + for version in HANDSHAKE_VERSIONS { + let resp = post_init(&client, &url, version.as_str()).await; + assert_eq!( + resp["result"]["protocolVersion"], + version.as_str(), + "supported version {version} should be echoed back" + ); + } + + ct.cancel(); +} + +#[tokio::test] +async fn stateless_json_init_does_not_agree_to_version_outside_supported_list() { + let (client, url, ct) = + spawn_server_of::(stateless_json_config()).await; + + let resp = post_init(&client, &url, ProtocolVersion::V_2026_07_28.as_str()).await; + assert_eq!( + resp["result"]["protocolVersion"], + ProtocolVersion::V_2025_11_25.as_str(), + "a version outside supported_protocol_versions should not be echoed back" ); ct.cancel(); diff --git a/crates/rmcp/tests/test_stateless_server_requests.rs b/crates/rmcp/tests/test_stateless_server_requests.rs new file mode 100644 index 0000000..9ab302b --- /dev/null +++ b/crates/rmcp/tests/test_stateless_server_requests.rs @@ -0,0 +1,288 @@ +#![cfg(all(feature = "server", not(feature = "local")))] + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use rmcp::{ + ServerHandler, ServiceExt, + model::{ + ClientCapabilities, ClientJsonRpcMessage, ClientRequest, DiscoverRequest, + DiscoverRequestParams, ErrorCode, ErrorData, Implementation, ListToolsRequest, + ListToolsResult, NumberOrString, PaginatedRequestParams, ProgressNotificationParam, + ProgressToken, ProtocolVersion, RequestId, RequestMetaObject, ServerJsonRpcMessage, + ServerNotification, + }, + service::{MaybeSendFuture, RequestContext, RoleServer, ServerInitializeError}, + transport::{IntoTransport, Transport}, +}; + +#[derive(Clone, Default)] +struct StatelessServer; + +impl ServerHandler for StatelessServer {} + +fn complete_meta() -> RequestMetaObject { + complete_meta_for("stateless-client") +} + +fn complete_meta_for(client_name: &str) -> RequestMetaObject { + let mut meta = RequestMetaObject::new(); + meta.set_protocol_version(ProtocolVersion::V_2026_07_28); + meta.set_client_info(Implementation::new(client_name, "1.0.0")); + meta.set_client_capabilities(ClientCapabilities::default()); + meta +} + +fn list_tools_request(meta: RequestMetaObject) -> ClientJsonRpcMessage { + let mut request = ListToolsRequest { + method: Default::default(), + params: None, + extensions: Default::default(), + }; + request.extensions.insert(meta); + ClientJsonRpcMessage::request( + ClientRequest::ListToolsRequest(request), + RequestId::Number(1), + ) +} + +#[tokio::test] +async fn stateless_server_rejects_missing_metadata_on_every_request() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server_task = tokio::spawn(async move { + StatelessServer + .serve(server_transport) + .await + .expect("server should start") + }); + let mut client = IntoTransport::::into_transport(client_transport); + + let mut discover = DiscoverRequest::new(DiscoverRequestParams {}); + discover.extensions.insert(complete_meta()); + client + .send(ClientJsonRpcMessage::request( + ClientRequest::DiscoverRequest(discover), + RequestId::Number(1), + )) + .await + .expect("send discover"); + assert!(matches!( + client.receive().await, + Some(ServerJsonRpcMessage::Response(_)) + )); + + client + .send(ClientJsonRpcMessage::request( + ClientRequest::ListToolsRequest(ListToolsRequest { + method: Default::default(), + params: None, + extensions: Default::default(), + }), + RequestId::Number(2), + )) + .await + .expect("send list tools"); + let Some(ServerJsonRpcMessage::Error(error)) = client.receive().await else { + panic!("expected invalid params"); + }; + assert_eq!(error.error.code, ErrorCode::INVALID_PARAMS); + + server_task + .await + .expect("server task") + .cancel() + .await + .expect("cancel server"); +} + +#[derive(Clone)] +struct ContextServer { + seen_clients: Arc>>, +} + +impl ServerHandler for ContextServer { + fn list_tools( + &self, + _request: Option, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + let seen_clients = self.seen_clients.clone(); + async move { + seen_clients + .lock() + .expect("seen clients lock") + .push(context.client_info().expect("current client info").name); + Ok(ListToolsResult::default()) + } + } +} + +#[tokio::test] +async fn stateless_server_uses_each_requests_client_context() { + let seen_clients = Arc::new(Mutex::new(Vec::new())); + let handler = ContextServer { + seen_clients: seen_clients.clone(), + }; + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server_task = tokio::spawn(async move { + handler + .serve(server_transport) + .await + .expect("server should start") + }); + let mut client = IntoTransport::::into_transport(client_transport); + + client + .send(list_tools_request(complete_meta_for("first-client"))) + .await + .expect("send first request"); + assert!(matches!( + client.receive().await, + Some(ServerJsonRpcMessage::Response(_)) + )); + + let mut second = list_tools_request(complete_meta_for("second-client")); + if let ClientJsonRpcMessage::Request(request) = &mut second { + request.id = RequestId::Number(2); + } + client.send(second).await.expect("send second request"); + assert!(matches!( + client.receive().await, + Some(ServerJsonRpcMessage::Response(_)) + )); + + assert_eq!( + *seen_clients.lock().expect("seen clients lock"), + ["first-client", "second-client"] + ); + server_task + .await + .expect("server task") + .cancel() + .await + .expect("cancel server"); +} + +#[tokio::test] +async fn stateless_server_rejects_malformed_metadata_opener_with_error_response() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server_task = tokio::spawn(async move { StatelessServer.serve(server_transport).await }); + let mut client = IntoTransport::::into_transport(client_transport); + + let mut request = ListToolsRequest { + method: Default::default(), + params: None, + extensions: Default::default(), + }; + let malformed: RequestMetaObject = serde_json::from_value(serde_json::json!({ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": "wrong", + "io.modelcontextprotocol/clientCapabilities": null + })) + .unwrap(); + request.extensions.insert(malformed); + client + .send(ClientJsonRpcMessage::request( + ClientRequest::ListToolsRequest(request), + RequestId::Number(1), + )) + .await + .expect("send list tools"); + let Some(ServerJsonRpcMessage::Error(error)) = client.receive().await else { + panic!("expected invalid params"); + }; + assert_eq!(error.id, Some(RequestId::Number(1))); + assert_eq!(error.error.code, ErrorCode::INVALID_PARAMS); + assert!( + error + .error + .message + .contains("request _meta is missing or has malformed required fields") + ); + assert!( + error + .error + .message + .contains("io.modelcontextprotocol/clientCapabilities") + ); + + let Err(error) = server_task.await.expect("server task") else { + panic!("malformed opener should not start a session"); + }; + assert!(matches!( + error, + ServerInitializeError::ExpectedInitializeRequest(Some(_)) + )); +} + +#[derive(Clone)] +struct ProgressServer; + +impl ServerHandler for ProgressServer { + async fn list_tools( + &self, + _request: Option, + context: RequestContext, + ) -> Result { + let progress_token = context + .meta + .get_progress_token() + .expect("progress token in request meta"); + context + .peer + .notify_progress(ProgressNotificationParam::new(progress_token, 1.0)) + .await + .expect("send progress notification"); + Ok(ListToolsResult::default()) + } +} + +/// Regression test for issue #1261: the first request of an `initialize`-less +/// session used to be handled outside the service loop, so a handler that sent +/// a notification before returning waited forever for the loop to flush it. +#[tokio::test] +async fn stateless_server_first_request_handler_can_send_notifications() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server_task = tokio::spawn(async move { + ProgressServer + .serve(server_transport) + .await + .expect("server should start") + }); + let mut client = IntoTransport::::into_transport(client_transport); + + let mut meta = complete_meta(); + meta.set_progress_token(ProgressToken(NumberOrString::Number(7))); + client + .send(list_tools_request(meta)) + .await + .expect("send first request"); + + let exchange = async { + let notification = match client.receive().await { + Some(ServerJsonRpcMessage::Notification(notification)) => notification.notification, + other => panic!("expected progress notification before the response, got {other:?}"), + }; + assert!( + matches!(notification, ServerNotification::ProgressNotification(_)), + "expected progress notification, got {notification:?}" + ); + let response = match client.receive().await { + Some(ServerJsonRpcMessage::Response(response)) => response, + other => panic!("expected list tools response, got {other:?}"), + }; + assert_eq!(response.id, RequestId::Number(1)); + }; + tokio::time::timeout(Duration::from_secs(5), exchange) + .await + .expect("first request must not deadlock the server"); + + server_task + .await + .expect("server task") + .cancel() + .await + .expect("cancel server"); +} diff --git a/crates/rmcp/tests/test_stdio_response_concurrency.rs b/crates/rmcp/tests/test_stdio_response_concurrency.rs index 563b0b3..64be3fc 100644 --- a/crates/rmcp/tests/test_stdio_response_concurrency.rs +++ b/crates/rmcp/tests/test_stdio_response_concurrency.rs @@ -4,7 +4,10 @@ use std::{collections::BTreeSet, process::Stdio, time::Duration}; use rmcp::{ ErrorData as McpError, ServerHandler, ServiceExt, - model::{CallToolRequestParams, CallToolResult, ContentBlock, ServerCapabilities, ServerInfo}, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ServerCapabilities, + ServerConfig, + }, }; use serde_json::{Value, json}; use tokio::{ @@ -79,19 +82,17 @@ struct LargeResponseServer; impl ServerHandler for LargeResponseServer { #[allow(deprecated)] - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) } async fn call_tool( &self, request: CallToolRequestParams, _context: rmcp::service::RequestContext, - ) -> Result { + ) -> Result { assert_eq!("large-response", request.name.as_ref()); - Ok(CallToolResult::success(vec![ContentBlock::text( - "x".repeat(RESPONSE_BYTES), - )])) + Ok(CallToolResult::success(vec![ContentBlock::text("x".repeat(RESPONSE_BYTES))]).into()) } } diff --git a/crates/rmcp/tests/test_streamable_http_4xx_error_body.rs b/crates/rmcp/tests/test_streamable_http_4xx_error_body.rs index ea49a41..13deeeb 100644 --- a/crates/rmcp/tests/test_streamable_http_4xx_error_body.rs +++ b/crates/rmcp/tests/test_streamable_http_4xx_error_body.rs @@ -7,7 +7,10 @@ use std::{collections::HashMap, sync::Arc}; use rmcp::{ - model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId}, + model::{ + ClientJsonRpcMessage, ClientRequest, DiscoverRequest, DiscoverRequestParams, PingRequest, + RequestId, + }, transport::streamable_http_client::{ StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse, }, @@ -45,6 +48,13 @@ fn ping_message() -> ClientJsonRpcMessage { ) } +fn discover_message() -> ClientJsonRpcMessage { + ClientJsonRpcMessage::request( + ClientRequest::DiscoverRequest(DiscoverRequest::new(DiscoverRequestParams {})), + RequestId::Number(1), + ) +} + /// HTTP 4xx with Content-Type: application/json and a valid JSON-RPC error body /// must be surfaced as `StreamableHttpPostResponse::Json`, not swallowed as a /// transport error. @@ -119,3 +129,26 @@ async fn http_4xx_malformed_json_body_falls_back_to_unexpected_server_response() other => panic!("expected UnexpectedServerResponse, got: {other:?}"), } } + +/// A discovery request attached to an existing session must retain the +/// transport error path; legacy fallback only applies to sessionless probes. +#[tokio::test] +async fn sessionful_discover_4xx_does_not_trigger_legacy_fallback() { + let url = spawn_mock_server(422, "text/plain", "Session-bound rejection").await; + + let client = reqwest::Client::new(); + let result = client + .post_message( + Arc::from(url.as_str()), + discover_message(), + Some(Arc::from("session-1")), + None, + HashMap::new(), + ) + .await; + + match result { + Err(StreamableHttpError::UnexpectedServerResponse(_)) => {} + other => panic!("expected UnexpectedServerResponse, got: {other:?}"), + } +} diff --git a/crates/rmcp/tests/test_streamable_http_client_concurrency.rs b/crates/rmcp/tests/test_streamable_http_client_concurrency.rs new file mode 100644 index 0000000..2e5be84 --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_client_concurrency.rs @@ -0,0 +1,1143 @@ +//! Independent http POSTs may overlap. A POST that reports an expired session +//! is retried at most once. +#![cfg(not(feature = "local"))] + +use std::{ + collections::HashMap, + io, + sync::{ + Arc, Mutex as StdMutex, + atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst}, + }, + time::Duration, +}; + +use futures::{StreamExt, stream::BoxStream}; +use http::{HeaderName, HeaderValue}; +use rmcp::{ + model::{ + CallToolRequestParams, CancelledNotificationParam, ClientConfig, ClientJsonRpcMessage, + ClientRequest, DiscoverResult, InitializeRequestParams, ProtocolVersion, Request, + RequestId, RequestMetaObject, ServerJsonRpcMessage, + }, + service::{ + ClientLifecycleMode, PeerRequestOptions, RequestHandle, RoleClient, RunningService, + serve_client_with_lifecycle, + }, + transport::streamable_http_client::{ + StreamableHttpClient, StreamableHttpClientTransport, StreamableHttpClientTransportConfig, + StreamableHttpError, StreamableHttpPostResponse, + }, +}; +use serde_json::{Value, json}; +use sse_stream::{Error as SseError, Sse}; +use tokio::{ + sync::{Mutex, mpsc, oneshot}, + task::JoinHandle, + time::timeout, +}; +use tokio_stream::wrappers::UnboundedReceiverStream; +use tokio_util::sync::CancellationToken; + +const TEST_TIMEOUT: Duration = Duration::from_secs(5); +type PostResult = Result>; +type Call = JoinHandle>; +type SseReceiver = mpsc::UnboundedReceiver>; + +#[derive(Default)] +struct Counts { + initialized: AtomicUsize, + hold_reinitialization: AtomicBool, + manual_controls: AtomicBool, + reject_unmatched_cancellation: AtomicBool, + local_streams: StdMutex>, + local_cancellations: AtomicUsize, + deleted: AtomicUsize, + cancelled: AtomicUsize, + posted: AtomicUsize, + active: AtomicUsize, + peak: AtomicUsize, +} + +struct ActivePost(Arc); + +impl Drop for ActivePost { + fn drop(&mut self) { + self.0.active.fetch_sub(1, SeqCst); + } +} + +struct LocalStreamRegistration { + id: RequestId, + counts: Arc, + dropped: Option>, +} + +impl Drop for LocalStreamRegistration { + fn drop(&mut self) { + self.counts.local_streams.lock().unwrap().remove(&self.id); + let _ = self.dropped.take().unwrap().send(()); + } +} + +struct Posted { + id: Value, + name: String, + session: Option>, + reply: oneshot::Sender, + returned: oneshot::Receiver<()>, +} + +struct ControlPost { + message: Value, + session: Option>, + reply: oneshot::Sender, +} + +fn response(id: Value, result: Value) -> ServerJsonRpcMessage { + serde_json::from_value(json!({ "jsonrpc": "2.0", "id": id, "result": result })) + .expect("valid scripted response") +} + +fn sse(message: Value) -> Result { + Ok(Sse { + event: Some("message".into()), + data: Some(message.to_string()), + id: None, + retry: None, + }) +} + +async fn next_event(receiver: &mut mpsc::UnboundedReceiver) -> T { + timeout(TEST_TIMEOUT, receiver.recv()) + .await + .expect("expected scripted event") + .expect("scripted client remains connected") +} + +impl Posted { + fn result(&self) -> ServerJsonRpcMessage { + response( + self.id.clone(), + json!({ "content": [{ "type": "text", "text": self.name }] }), + ) + } + + fn finish(self, result: PostResult) { + self.reply.send(result).expect("POST is still waiting"); + } + + fn succeed(self) { + let result = StreamableHttpPostResponse::Json(self.result(), None); + self.finish(Ok(result)); + } + + fn expire(self) { + self.finish(Err(StreamableHttpError::SessionExpired)); + } + + async fn finish_and_wait(self, result: PostResult) -> anyhow::Result<()> { + let Self { + reply, returned, .. + } = self; + reply.send(result).expect("POST is still waiting"); + timeout(TEST_TIMEOUT, returned).await??; + Ok(()) + } + + async fn expire_and_wait(self) -> anyhow::Result<()> { + self.finish_and_wait(Err(StreamableHttpError::SessionExpired)) + .await + } + + async fn start_sse(self) -> anyhow::Result> { + let message = serde_json::to_value(self.result()).unwrap(); + let (release, released) = oneshot::channel(); + let stream = futures::stream::once(async move { + released.await.expect("release the SSE response"); + sse(message) + }) + .boxed(); + self.finish_and_wait(Ok(StreamableHttpPostResponse::Sse(stream, None))) + .await?; + Ok(release) + } + + async fn start_local_sse( + self, + counts: Arc, + ) -> anyhow::Result<(CancellationToken, oneshot::Receiver<()>)> { + let id: RequestId = serde_json::from_value(self.id.clone())?; + let cancellation = CancellationToken::new(); + let eof = cancellation.clone(); + counts + .local_streams + .lock() + .unwrap() + .insert(id.clone(), cancellation.clone()); + let (dropped, closed) = oneshot::channel(); + let registration = LocalStreamRegistration { + id, + counts, + dropped: Some(dropped), + }; + let (started, listening) = oneshot::channel(); + let stream = futures::stream::once(async move { + let _registration = registration; + let _ = started.send(()); + cancellation.cancelled().await; + None::> + }) + .filter_map(futures::future::ready) + .boxed(); + self.finish_and_wait(Ok(StreamableHttpPostResponse::Sse(stream, None))) + .await?; + timeout(TEST_TIMEOUT, listening).await??; + Ok((eof, closed)) + } +} + +#[derive(Clone)] +struct ScriptedClient { + started: mpsc::UnboundedSender, + controls: mpsc::UnboundedSender, + incoming: Arc>>, + reinitializing: mpsc::UnboundedSender>, + counts: Arc, +} + +impl ScriptedClient { + async fn control_post(&self, message: Value, session: Option>) -> PostResult { + if !self.counts.manual_controls.load(SeqCst) { + return Ok(StreamableHttpPostResponse::Accepted); + } + let (reply, result) = oneshot::channel(); + self.controls + .send(ControlPost { + message, + session, + reply, + }) + .expect("test remains connected"); + result.await.expect("test answers the control POST") + } +} + +impl StreamableHttpClient for ScriptedClient { + type Error = io::Error; + + async fn post_message( + &self, + _uri: Arc, + message: ClientJsonRpcMessage, + session: Option>, + _auth_header: Option, + _custom_headers: HashMap, + ) -> PostResult { + let value = serde_json::to_value(message).unwrap(); + match value["method"].as_str() { + Some("server/discover") => Ok(StreamableHttpPostResponse::Json( + response( + value["id"].clone(), + serde_json::to_value(DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + serde_json::from_value(json!({ "tools": {} })).unwrap(), + )) + .unwrap(), + ), + None, + )), + Some("initialize") => { + let generation = self.counts.initialized.fetch_add(1, SeqCst) + 1; + if generation > 1 && self.counts.hold_reinitialization.load(SeqCst) { + let (release, released) = oneshot::channel(); + self.reinitializing + .send(release) + .expect("test remains connected"); + released.await.expect("test releases reinitialization"); + } + Ok(StreamableHttpPostResponse::Json( + response( + value["id"].clone(), + json!({ + "protocolVersion": "2025-11-25", + "capabilities": { "tools": {} }, + "serverInfo": { "name": "scripted", "version": "1" }, + }), + ), + Some(format!("session-{generation}")), + )) + } + Some("notifications/initialized") => Ok(StreamableHttpPostResponse::Accepted), + Some("notifications/cancelled") => { + let id: RequestId = + serde_json::from_value(value["params"]["requestId"].clone()).unwrap(); + let local = self.counts.local_streams.lock().unwrap().get(&id).cloned(); + if let Some(cancellation) = local { + self.counts.local_cancellations.fetch_add(1, SeqCst); + cancellation.cancel(); + return Ok(StreamableHttpPostResponse::Accepted); + } + self.counts.cancelled.fetch_add(1, SeqCst); + if self.counts.reject_unmatched_cancellation.load(SeqCst) { + return Err(StreamableHttpError::Client(io::Error::other( + "local event cancellation reached the http POST path", + ))); + } + self.control_post(value, session).await + } + Some("tools/call") => { + self.counts.posted.fetch_add(1, SeqCst); + let active = self.counts.active.fetch_add(1, SeqCst) + 1; + self.counts.peak.fetch_max(active, SeqCst); + let _active = ActivePost(self.counts.clone()); + let (reply, response) = oneshot::channel(); + let (finished, returned) = oneshot::channel(); + self.started + .send(Posted { + id: value["id"].clone(), + name: value["params"]["name"].as_str().unwrap().to_owned(), + session, + reply, + returned, + }) + .expect("test remains connected"); + let response = response.await.expect("test answers each POST"); + let _ = finished.send(()); + response + } + None if value.get("result").is_some() || value.get("error").is_some() => { + self.control_post(value, session).await + } + method => panic!("unexpected scripted method: {method:?}"), + } + } + + async fn delete_session( + &self, + _uri: Arc, + _session: Arc, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + assert_eq!( + self.counts.active.load(SeqCst), + 0, + "POSTs must stop before deleting the session" + ); + self.counts.deleted.fetch_add(1, SeqCst); + Ok(()) + } + + async fn get_stream( + &self, + _uri: Arc, + _session: Option>, + _last_event_id: Option, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result>, StreamableHttpError> { + Ok(match self.incoming.lock().await.take() { + Some(incoming) => UnboundedReceiverStream::new(incoming).boxed(), + None => futures::stream::pending().boxed(), + }) + } +} + +struct Harness { + client: RunningService, + started: mpsc::UnboundedReceiver, + controls: mpsc::UnboundedReceiver, + incoming: mpsc::UnboundedSender>, + reinitializations: mpsc::UnboundedReceiver>, + counts: Arc, +} + +fn config() -> StreamableHttpClientTransportConfig { + StreamableHttpClientTransportConfig::with_uri("http://scripted/mcp") +} + +fn transport_error(error: &anyhow::Error) -> &StreamableHttpError { + let service_error = error + .downcast_ref::() + .expect("expected a service error"); + let rmcp::ServiceError::TransportSend(transport_error) = service_error else { + panic!("expected a transport error, got {service_error:?}"); + }; + transport_error + .error + .downcast_ref::>() + .expect("expected a streamable http error") +} + +fn assert_recovery_timeout(error: anyhow::Error) { + assert!(matches!( + transport_error(&error), + StreamableHttpError::SessionRecoveryTimeout + )); +} + +impl Harness { + async fn start(config: StreamableHttpClientTransportConfig) -> anyhow::Result { + Self::with_lifecycle(config, ClientLifecycleMode::Initialize).await + } + + async fn with_lifecycle( + config: StreamableHttpClientTransportConfig, + lifecycle: ClientLifecycleMode, + ) -> anyhow::Result { + let (started, requests) = mpsc::unbounded_channel(); + let (control_tx, controls) = mpsc::unbounded_channel(); + let (incoming, incoming_rx) = mpsc::unbounded_channel(); + let (reinitializing, reinitializations) = mpsc::unbounded_channel(); + let counts = Arc::new(Counts::default()); + let transport = StreamableHttpClientTransport::with_client( + ScriptedClient { + started, + controls: control_tx, + incoming: Arc::new(Mutex::new(Some(incoming_rx))), + reinitializing, + counts: counts.clone(), + }, + config, + ); + let client = + serve_client_with_lifecycle(ClientConfig::default(), transport, lifecycle).await?; + Ok(Self { + client, + started: requests, + controls, + incoming, + reinitializations, + counts, + }) + } + + fn call(&self, name: impl Into) -> Call { + let name = name.into(); + let peer = self.client.peer().clone(); + tokio::spawn(async move { + let result = peer + .call_tool(CallToolRequestParams::new(name.clone())) + .await?; + anyhow::ensure!(serde_json::to_value(result)?["content"][0]["text"] == name); + Ok(()) + }) + } + + async fn cancellable(&self, name: &'static str) -> anyhow::Result> { + self.cancellable_with_options(name, PeerRequestOptions::no_options()) + .await + } + + async fn cancellable_with_options( + &self, + name: &'static str, + options: PeerRequestOptions, + ) -> anyhow::Result> { + Ok(self + .client + .peer() + .send_cancellable_request( + ClientRequest::CallToolRequest(Request::new(CallToolRequestParams::new(name))), + options, + ) + .await?) + } + + fn notify_cancellation(&self, id: RequestId) -> JoinHandle> { + let peer = self.client.peer().clone(); + tokio::spawn(async move { + peer.notify_cancelled(CancelledNotificationParam::new(Some(id), None)) + .await + }) + } + + async fn next(&mut self) -> Posted { + next_event(&mut self.started).await + } + + async fn next_control(&mut self) -> ControlPost { + next_event(&mut self.controls).await + } + + async fn exchange_ping(&mut self, id: &str) { + self.incoming + .send(sse(json!({ "jsonrpc": "2.0", "id": id, "method": "ping" }))) + .expect("common SSE stream remains open"); + let control = self.next_control().await; + assert_eq!(control.message["id"], id); + assert!(control.message["result"].is_object()); + assert_eq!(control.session.as_deref(), Some("session-1")); + control + .reply + .send(Ok(StreamableHttpPostResponse::Accepted)) + .expect("reply POST is still waiting"); + } + + async fn finish( + self, + calls: Vec, + posted: usize, + initialized: usize, + ) -> anyhow::Result<()> { + for call in calls { + timeout(TEST_TIMEOUT, call).await???; + } + assert_eq!(self.counts.posted.load(SeqCst), posted); + assert_eq!(self.counts.initialized.load(SeqCst), initialized); + assert_eq!(self.counts.active.load(SeqCst), 0); + self.client.cancel().await?; + Ok(()) + } +} + +#[tokio::test] +async fn json_limits_allow_overlap_and_preserve_response_ids() -> anyhow::Result<()> { + let mut zero = config(); + zero.max_concurrent_requests = 0; + for (config, limit, total) in [ + (config().max_concurrent_requests(2), 2, 5), + (config().max_concurrent_requests(1), 1, 3), + (zero, 1, 2), + (config(), 16, 17), + ] { + let mut harness = Harness::start(config).await?; + let calls = (0..total) + .map(|index| harness.call(format!("request-{index}"))) + .collect(); + let mut pending = Vec::new(); + for _ in 0..limit { + pending.push(harness.next().await); + } + assert_eq!(harness.counts.active.load(SeqCst), limit); + // Keep the oldest response blocked while newer requests finish first. + for _ in limit..total { + pending.pop().unwrap().succeed(); + pending.push(harness.next().await); + } + for request in pending.into_iter().rev() { + request.succeed(); + } + let counts = harness.counts.clone(); + harness.finish(calls, total, 1).await?; + assert_eq!(counts.peak.load(SeqCst), limit); + } + Ok(()) +} + +#[tokio::test] +async fn early_sse_response_releases_the_post_slot() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(1)).await?; + let first = harness.call("first"); + let release = harness.next().await.start_sse().await?; + let second = harness.call("second"); + harness.next().await.succeed(); + timeout(TEST_TIMEOUT, second).await???; + assert!(!first.is_finished(), "the SSE response is still blocked"); + release.send(()).unwrap(); + assert_eq!(harness.counts.peak.load(SeqCst), 1); + harness.finish(vec![first], 2, 1).await +} + +#[tokio::test] +async fn a_common_stream_response_finishes_the_request_before_its_post_returns() +-> anyhow::Result<()> { + let (stream, incoming) = mpsc::unbounded_channel(); + for (late, orphan) in [ + (Ok(StreamableHttpPostResponse::Accepted), None), + (Err(StreamableHttpError::SessionExpired), None), + ( + Ok(StreamableHttpPostResponse::Sse( + UnboundedReceiverStream::new(incoming).boxed(), + None, + )), + Some(stream), + ), + ] { + let mut harness = Harness::start(config().max_concurrent_requests(1)).await?; + let request = harness.cancellable("completed-on-common-stream").await?; + let post = harness.next().await; + harness + .incoming + .send(sse(serde_json::to_value(post.result())?))?; + timeout(TEST_TIMEOUT, request.await_response()).await??; + + post.finish_and_wait(late).await?; + let next = harness.call("after-completed-request"); + let post = harness.next().await; + assert_eq!( + post.name, "after-completed-request", + "do not retry completed work" + ); + assert_eq!(post.session.as_deref(), Some("session-1")); + post.succeed(); + timeout(TEST_TIMEOUT, next).await???; + if let Some(stream) = orphan { + timeout(Duration::from_secs(1), stream.closed()) + .await + .expect("a late SSE response must not leave an orphan stream"); + } + harness.finish(vec![], 2, 1).await?; + } + Ok(()) +} + +#[tokio::test] +async fn common_stream_response_keeps_numeric_and_string_ids_distinct() -> anyhow::Result<()> { + use rmcp::transport::Transport; + + let (started, mut requests) = mpsc::unbounded_channel(); + let (incoming, incoming_rx) = mpsc::unbounded_channel(); + let counts = Arc::new(Counts::default()); + let mut transport = StreamableHttpClientTransport::with_client( + ScriptedClient { + started, + controls: mpsc::unbounded_channel().0, + incoming: Arc::new(Mutex::new(Some(incoming_rx))), + reinitializing: mpsc::unbounded_channel().0, + counts: counts.clone(), + }, + config(), + ); + for message in [ + json!({ + "jsonrpc": "2.0", + "id": 0, + "method": "initialize", + "params": InitializeRequestParams::default() + }), + json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }), + ] { + timeout( + TEST_TIMEOUT, + transport.send(serde_json::from_value(message)?), + ) + .await??; + } + assert!(timeout(TEST_TIMEOUT, transport.receive()).await?.is_some()); + let call = |id, name| { + ClientJsonRpcMessage::request( + ClientRequest::CallToolRequest(Request::new(CallToolRequestParams::new(name))), + id, + ) + }; + + let mut numeric_send = Box::pin(transport.send(call(RequestId::Number(7), "numeric"))); + assert!(futures::poll!(numeric_send.as_mut()).is_pending()); + let numeric_post = next_event(&mut requests).await; + let numeric_result = numeric_post.result(); + let release_numeric = numeric_post.start_sse().await?; + timeout(TEST_TIMEOUT, numeric_send).await??; + + // This exact string id is pending, but has no response stream registration. + let mut string_send = Box::pin(transport.send(call(RequestId::String("7".into()), "string"))); + assert!(futures::poll!(string_send.as_mut()).is_pending()); + let string_post = next_event(&mut requests).await; + let string_result = serde_json::to_value(string_post.result())?; + string_post.finish(Ok(StreamableHttpPostResponse::Accepted)); + timeout(TEST_TIMEOUT, string_send).await??; + + incoming.send(sse(string_result.clone()))?; + let received = timeout(TEST_TIMEOUT, transport.receive()) + .await? + .expect("string-id response"); + assert_eq!(serde_json::to_value(received)?, string_result); + + release_numeric + .send(()) + .expect("the distinct numeric-id stream must remain open"); + let received = timeout(TEST_TIMEOUT, transport.receive()) + .await? + .expect("numeric-id response"); + assert_eq!( + serde_json::to_value(received)?, + serde_json::to_value(numeric_result)? + ); + assert_eq!(counts.posted.load(SeqCst), 2); + assert_eq!(counts.cancelled.load(SeqCst), 0); + transport.close().await?; + Ok(()) +} + +#[tokio::test] +async fn a_common_stream_response_prevents_replay_during_session_recovery() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(2)).await?; + let blocking = harness.call("blocking-recovery"); + let blocked = harness.next().await; + let request = harness.cancellable("completed-during-recovery").await?; + let expired = harness.next().await; + let final_response = sse(serde_json::to_value(expired.result())?); + expired.expire_and_wait().await?; + harness.incoming.send(final_response)?; + timeout(TEST_TIMEOUT, request.await_response()).await??; + blocked.succeed(); + timeout(TEST_TIMEOUT, blocking).await???; + + let next = harness.call("after-completed-request"); + let post = harness.next().await; + assert_eq!( + post.name, "after-completed-request", + "do not retry completed work" + ); + let initializations = harness.counts.initialized.load(SeqCst); + assert!((1..=2).contains(&initializations)); + post.succeed(); + timeout(TEST_TIMEOUT, next).await???; + harness.finish(vec![], 3, initializations).await +} + +#[tokio::test] +async fn concurrent_session_expiry_shares_one_reinitialization() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(2)).await?; + let calls = vec![harness.call("first"), harness.call("second")]; + // Hold both requests before releasing either expired response. + let first = harness.next().await; + let second = harness.next().await; + let mut originals = HashMap::new(); + for request in [first, second] { + assert_eq!(request.session.as_deref(), Some("session-1")); + originals.insert(request.name.clone(), request.id.clone()); + request.expire(); + } + for _ in 0..2 { + let retry = harness.next().await; + assert_eq!(retry.session.as_deref(), Some("session-2")); + assert_eq!(originals.remove(&retry.name), Some(retry.id.clone())); + retry.succeed(); + } + assert!(originals.is_empty()); + harness.finish(calls, 4, 2).await +} + +#[tokio::test] +async fn cancellation_still_runs_while_recovery_waits_for_old_posts() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(3)).await?; + let hanging = harness.cancellable("hanging").await?; + let mut blocked = harness.next().await; + let expired = harness.call("expired"); + harness.next().await.expire_and_wait().await?; + assert_eq!(harness.counts.initialized.load(SeqCst), 1); + assert_eq!(harness.counts.active.load(SeqCst), 1); + + timeout(Duration::from_secs(1), hanging.cancel(None)) + .await + .expect("cancellation must bypass the session recovery wait")?; + timeout(TEST_TIMEOUT, blocked.reply.closed()).await?; + let retry = harness.next().await; + assert_eq!(retry.name, "expired"); + assert_eq!(retry.session.as_deref(), Some("session-2")); + retry.succeed(); + harness.finish(vec![expired], 3, 2).await +} + +#[tokio::test] +async fn server_replies_still_run_while_recovery_waits_for_old_posts() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(2)).await?; + harness.counts.manual_controls.store(true, SeqCst); + let waiting = harness.call("waiting-for-client"); + let blocked = harness.next().await; + let expired = harness.call("expired"); + harness.next().await.expire_and_wait().await?; + + harness.exchange_ping("recovery-ping").await; + blocked.succeed(); + let retry = harness.next().await; + assert_eq!(retry.name, "expired"); + assert_eq!(retry.session.as_deref(), Some("session-2")); + retry.succeed(); + harness.finish(vec![waiting, expired], 3, 2).await +} + +#[tokio::test] +async fn a_version_barrier_allows_the_server_reply_it_needs() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(2)).await?; + harness.counts.manual_controls.store(true, SeqCst); + let mut meta = RequestMetaObject::new(); + meta.set_protocol_version(ProtocolVersion::V_2025_06_18); + let barrier = harness + .cancellable_with_options("barrier", PeerRequestOptions::no_options().with_meta(meta)) + .await?; + let blocked = harness.next().await; + let queued = harness.cancellable("after-barrier").await?; + + harness.exchange_ping("barrier-ping").await; + assert_eq!(harness.counts.posted.load(SeqCst), 1); + blocked.succeed(); + timeout(TEST_TIMEOUT, barrier.await_response()).await??; + let next = harness.next().await; + assert_eq!(next.name, "after-barrier"); + next.succeed(); + timeout(TEST_TIMEOUT, queued.await_response()).await??; + harness.finish(vec![], 2, 1).await +} + +#[tokio::test] +async fn a_cancelled_version_barrier_does_not_block_unrelated_posts() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(2)).await?; + let blocking = harness.call("blocking"); + let blocked = harness.next().await; + let mut meta = RequestMetaObject::new(); + meta.set_protocol_version(ProtocolVersion::V_2025_06_18); + let barrier = harness + .cancellable_with_options( + "cancelled-barrier", + PeerRequestOptions::no_options().with_meta(meta), + ) + .await?; + timeout(TEST_TIMEOUT, barrier.cancel(None)).await??; + + let next = harness.call("after-cancelled-barrier"); + let post = harness.next().await; + assert_eq!(post.name, "after-cancelled-barrier"); + post.succeed(); + timeout(TEST_TIMEOUT, next).await???; + assert!(!blocked.reply.is_closed(), "the older POST is still held"); + blocked.succeed(); + harness.finish(vec![blocking], 2, 1).await +} + +#[tokio::test] +async fn dropping_a_parked_barrier_send_unblocks_the_ordinary_queue() -> anyhow::Result<()> { + use std::task::{Context, Wake, Waker}; + + use rmcp::transport::Transport; + + struct WakeSignal(mpsc::UnboundedSender<()>); + impl Wake for WakeSignal { + fn wake(self: Arc) { + let _ = self.0.send(()); + } + } + + let (started, mut requests) = mpsc::unbounded_channel(); + let counts = Arc::new(Counts::default()); + let mut config = config().max_concurrent_requests(2); + config.channel_buffer_capacity = 1; + let mut transport = StreamableHttpClientTransport::with_client( + ScriptedClient { + started, + controls: mpsc::unbounded_channel().0, + incoming: Arc::new(Mutex::new(None)), + reinitializing: mpsc::unbounded_channel().0, + counts: counts.clone(), + }, + config, + ); + for message in [ + json!({ + "jsonrpc": "2.0", + "id": 0, + "method": "initialize", + "params": InitializeRequestParams::default() + }), + json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }), + ] { + timeout( + TEST_TIMEOUT, + transport.send(serde_json::from_value(message)?), + ) + .await??; + } + assert!(timeout(TEST_TIMEOUT, transport.receive()).await?.is_some()); + let call = |id, params| { + ClientJsonRpcMessage::request( + ClientRequest::CallToolRequest(Request::new(params)), + RequestId::Number(id), + ) + }; + let mut first = Box::pin(transport.send(call(1, CallToolRequestParams::new("held")))); + assert!(futures::poll!(first.as_mut()).is_pending()); + let blocked = next_event(&mut requests).await; + + let mut meta = RequestMetaObject::new(); + meta.set_protocol_version(ProtocolVersion::V_2025_06_18); + let mut barrier_request = Request::new(CallToolRequestParams::new("abandoned-barrier")); + barrier_request.extensions.insert(meta); + let mut barrier = Box::pin(transport.send(ClientJsonRpcMessage::request( + ClientRequest::CallToolRequest(barrier_request), + RequestId::Number(2), + ))); + assert!(futures::poll!(barrier.as_mut()).is_pending()); + let mut next = Box::pin(transport.send(call(3, CallToolRequestParams::new("after-barrier")))); + let (wake, mut woke) = mpsc::unbounded_channel(); + let waker = Waker::from(Arc::new(WakeSignal(wake))); + assert!( + next.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + // The barrier fills the queue. The next send wakes only after the worker parks it. + next_event(&mut woke).await; + assert_eq!(counts.posted.load(SeqCst), 1, "the barrier must be parked"); + drop(barrier); + assert!(futures::poll!(next.as_mut()).is_pending()); + + let post = timeout(Duration::from_secs(1), requests.recv()) + .await + .expect("dropping the parked send must wake the worker") + .unwrap(); + assert_eq!(post.name, "after-barrier"); + post.succeed(); + timeout(TEST_TIMEOUT, next).await??; + assert!(timeout(TEST_TIMEOUT, transport.receive()).await?.is_some()); + assert!(!blocked.reply.is_closed(), "the first POST remains held"); + blocked.succeed(); + timeout(TEST_TIMEOUT, first).await??; + assert_eq!(counts.posted.load(SeqCst), 2); + assert_eq!(counts.cancelled.load(SeqCst), 0); + transport.close().await?; + Ok(()) +} + +#[tokio::test] +async fn recovery_deadline_drops_ambiguous_posts_without_retrying_them() -> anyhow::Result<()> { + assert_eq!(config().session_recovery_timeout, Duration::from_secs(5)); + let mut harness = Harness::start( + config() + .max_concurrent_requests(3) + .session_recovery_timeout(Duration::from_millis(50)), + ) + .await?; + let ambiguous = harness.call("possibly-applied"); + let mut blocked = harness.next().await; + let expired = harness.call("expired"); + let rejected = harness.next().await; + let rejected_id = rejected.id.clone(); + rejected.expire_and_wait().await?; + + assert_recovery_timeout(timeout(TEST_TIMEOUT, ambiguous).await??.unwrap_err()); + timeout(TEST_TIMEOUT, blocked.reply.closed()).await?; + let retry = harness.next().await; + assert_eq!(retry.name, "expired"); + assert_eq!(retry.id, rejected_id); + assert_eq!(retry.session.as_deref(), Some("session-2")); + retry.succeed(); + harness.finish(vec![expired], 3, 2).await +} + +#[tokio::test] +async fn reinitialization_has_its_own_deadline() -> anyhow::Result<()> { + let mut harness = + Harness::start(config().session_recovery_timeout(Duration::from_millis(50))).await?; + harness.counts.hold_reinitialization.store(true, SeqCst); + let expired = harness.call("expired"); + harness.next().await.expire_and_wait().await?; + let mut reinitialization = next_event(&mut harness.reinitializations).await; + + assert_recovery_timeout(timeout(TEST_TIMEOUT, expired).await??.unwrap_err()); + timeout(TEST_TIMEOUT, reinitialization.closed()).await?; + harness.finish(vec![], 1, 2).await +} + +#[tokio::test] +async fn an_expired_retry_is_not_retried_again() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(2)).await?; + let call = harness.call("expires-twice"); + let first = harness.next().await; + let id = first.id.clone(); + first.expire(); + let retry = harness.next().await; + assert_eq!(retry.id, id); + assert_eq!(retry.session.as_deref(), Some("session-2")); + retry.expire(); + let error = timeout(TEST_TIMEOUT, call).await??.unwrap_err(); + assert!(error.to_string().contains("Session expired")); + harness.finish(vec![], 2, 2).await +} + +#[tokio::test] +async fn a_lost_post_response_is_not_retried() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(2)).await?; + let call = harness.call("possibly-applied"); + harness + .next() + .await + .finish(Err(StreamableHttpError::Client(io::Error::new( + io::ErrorKind::ConnectionReset, + "scripted response lost", + )))); + let error = timeout(TEST_TIMEOUT, call).await??.unwrap_err(); + assert!(error.to_string().contains("scripted response lost")); + harness.finish(vec![], 1, 1).await +} + +#[tokio::test] +async fn cancellation_bypasses_queued_posts_at_capacity() -> anyhow::Result<()> { + for (lifecycle, legacy_notifications, initializations) in [ + (ClientLifecycleMode::Initialize, 2, 1), + ( + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + 0, + 0, + ), + ] { + let mut harness = + Harness::with_lifecycle(config().max_concurrent_requests(1), lifecycle).await?; + let request = harness.cancellable("cancel-me").await?; + let mut blocked = harness.next().await; + let cancelled = harness.cancellable("never-send").await?; + timeout(TEST_TIMEOUT, cancelled.cancel(None)).await??; + let queued = harness.cancellable("queued").await?; + timeout(TEST_TIMEOUT, request.cancel(None)).await??; + timeout(TEST_TIMEOUT, blocked.reply.closed()).await?; + let next = harness.next().await; + assert_eq!(next.name, "queued"); + next.succeed(); + timeout(TEST_TIMEOUT, queued.await_response()).await??; + assert_eq!(harness.counts.cancelled.load(SeqCst), legacy_notifications); + harness.finish(vec![], 2, initializations).await?; + } + Ok(()) +} + +#[tokio::test] +async fn legacy_cancellation_reaches_the_adapter_before_its_local_stream_drops() +-> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(1)).await?; + harness.counts.manual_controls.store(true, SeqCst); + let stale = harness.notify_cancellation(RequestId::Number(999)); + let held = harness.next_control().await; + harness.counts.manual_controls.store(false, SeqCst); + harness + .counts + .reject_unmatched_cancellation + .store(true, SeqCst); + + let event = harness.cancellable("local-events").await?; + let (eof, dropped) = harness + .next() + .await + .start_local_sse(harness.counts.clone()) + .await?; + let peer = harness.client.peer().clone(); + let mut cancel = Box::pin(peer.notify_cancelled(CancelledNotificationParam::new( + Some(event.id.clone()), + None, + ))); + assert!(futures::poll!(cancel.as_mut()).is_pending()); + // The ordinary request follows the cancellation through the service queue. + let probe = harness.call("after-cancellation"); + harness.next().await.succeed(); + timeout(TEST_TIMEOUT, probe).await???; + + // EOF is ready, but queued cancellation must pause reads until dispatch. + eof.cancel(); + let probe = harness.call("after-eof"); + harness.next().await.succeed(); + timeout(TEST_TIMEOUT, probe).await???; + assert!( + harness + .counts + .local_streams + .lock() + .unwrap() + .contains_key(&event.id), + "the adapter's stream registration must survive queued cancellation" + ); + assert!(!held.reply.is_closed(), "the control slot is still held"); + + held.reply + .send(Ok(StreamableHttpPostResponse::Accepted)) + .unwrap(); + timeout(TEST_TIMEOUT, stale).await???; + timeout(TEST_TIMEOUT, cancel).await??; + timeout(TEST_TIMEOUT, dropped).await??; + assert!(matches!( + timeout(TEST_TIMEOUT, event.await_response()).await?, + Err(rmcp::ServiceError::Cancelled { .. }) + )); + assert_eq!(harness.counts.local_cancellations.load(SeqCst), 1); + assert_eq!(harness.counts.cancelled.load(SeqCst), 1); + harness.finish(vec![], 3, 1).await +} + +#[tokio::test] +async fn control_timeout_is_configurable_and_releases_its_slot() -> anyhow::Result<()> { + assert_eq!(config().control_request_timeout, Duration::from_secs(5)); + let mut harness = + Harness::start(config().control_request_timeout(Duration::from_millis(50))).await?; + harness.counts.manual_controls.store(true, SeqCst); + let cancellation = harness.notify_cancellation(RequestId::Number(999)); + let mut cancelled_post = harness.next_control().await; + harness.incoming.send(sse(json!({ + "jsonrpc": "2.0", "id": "timed-reply", "method": "ping" + })))?; + + let error = timeout(Duration::from_secs(1), cancellation) + .await?? + .unwrap_err(); + assert!(matches!( + transport_error(&error.into()), + StreamableHttpError::ControlRequestTimeout + )); + timeout(Duration::from_secs(1), cancelled_post.reply.closed()).await?; + + let mut reply_post = harness.next_control().await; + assert_eq!(reply_post.message["id"], "timed-reply"); + assert!(reply_post.message["result"].is_object()); + harness.counts.manual_controls.store(false, SeqCst); + let next = harness.notify_cancellation(RequestId::Number(1000)); + timeout(Duration::from_secs(1), reply_post.reply.closed()).await?; + timeout(Duration::from_secs(1), next).await???; + assert_eq!(harness.counts.cancelled.load(SeqCst), 2); + harness.finish(vec![], 0, 1).await +} + +#[tokio::test] +async fn a_hanging_legacy_control_does_not_delay_local_cancellation() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(1)).await?; + harness.counts.manual_controls.store(true, SeqCst); + let stale = harness.notify_cancellation(RequestId::Number(999)); + let mut held = harness.next_control().await; + assert_eq!(held.message["params"]["requestId"], 999); + harness.counts.manual_controls.store(false, SeqCst); + + let live = harness.cancellable("live-post").await?; + let mut blocked = harness.next().await; + let cancel_post = tokio::spawn(async move { live.cancel(None).await }); + timeout(Duration::from_secs(1), blocked.reply.closed()).await?; + + let streaming = harness.cancellable("live-stream").await?; + let mut stream = harness.next().await.start_sse().await?; + let cancel_stream = harness.notify_cancellation(streaming.id.clone()); + assert!( + !held.reply.is_closed(), + "the old control POST is still held" + ); + assert_eq!(harness.counts.cancelled.load(SeqCst), 1); + + // The default control timeout is five seconds; give its watchdog headroom. + let error = timeout(Duration::from_secs(10), stale).await??.unwrap_err(); + assert!(matches!( + transport_error(&error.into()), + StreamableHttpError::ControlRequestTimeout + )); + timeout(TEST_TIMEOUT, held.reply.closed()).await?; + timeout(TEST_TIMEOUT, cancel_post).await???; + timeout(TEST_TIMEOUT, cancel_stream).await???; + timeout(TEST_TIMEOUT, stream.closed()).await?; + assert!(matches!( + timeout(TEST_TIMEOUT, streaming.await_response()).await?, + Err(rmcp::ServiceError::Cancelled { .. }) + )); + assert_eq!(harness.counts.cancelled.load(SeqCst), 3); + harness.finish(vec![], 2, 1).await +} + +#[tokio::test] +async fn close_drops_blocked_posts_before_deleting_the_session() -> anyhow::Result<()> { + let mut harness = Harness::start(config().max_concurrent_requests(2)).await?; + let calls = [harness.call("first"), harness.call("second")]; + let posts = [harness.next().await, harness.next().await]; + timeout(TEST_TIMEOUT, harness.client.cancel()).await??; + assert!(posts.iter().all(|post| post.reply.is_closed())); + assert_eq!(harness.counts.active.load(SeqCst), 0); + assert_eq!(harness.counts.deleted.load(SeqCst), 1); + for call in calls { + assert!(timeout(TEST_TIMEOUT, call).await??.is_err()); + } + Ok(()) +} diff --git a/crates/rmcp/tests/test_streamable_http_connection_reuse.rs b/crates/rmcp/tests/test_streamable_http_connection_reuse.rs index 553448e..aa4ae63 100644 --- a/crates/rmcp/tests/test_streamable_http_connection_reuse.rs +++ b/crates/rmcp/tests/test_streamable_http_connection_reuse.rs @@ -5,7 +5,7 @@ use std::time::Instant; use rmcp::{ ServerHandler, ServiceExt, handler::server::{router::tool::ToolRouter, wrapper::Parameters}, - model::{CallToolRequestParams, ClientInfo, ServerCapabilities, ServerInfo}, + model::{CallToolRequestParams, ClientConfig, ServerCapabilities, ServerConfig}, schemars, tool, tool_handler, tool_router, transport::{ StreamableHttpClientTransport, @@ -46,8 +46,8 @@ impl SumServer { #[tool_handler(router = self.tool_router)] impl ServerHandler for SumServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) } } @@ -83,7 +83,7 @@ async fn test_subsequent_tool_calls_reuse_connections() -> anyhow::Result<()> { let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(format!("http://{addr}/mcp")), ); - let client = ClientInfo::default().serve(transport).await?; + let client = ClientConfig::default().serve(transport).await?; // Warm up: first call may include one-time setup costs. let args: serde_json::Map = diff --git a/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs b/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs new file mode 100644 index 0000000..05efe21 --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs @@ -0,0 +1,196 @@ +#![cfg(all( + feature = "server", + feature = "transport-streamable-http-server", + feature = "reqwest", + not(feature = "local") +))] + +//! Regression test for #857: when a stateless streamable-HTTP client disconnects +//! (drops the response) while a tool handler is still awaiting, the per-request +//! `RequestContext::ct` should fire so the handler can cancel cooperatively. +//! +//! Stateless requests are one-shot (no session, no resumption), so a dropped +//! response is terminal and safe to cancel — unlike the stateful/resumable path, +//! where a disconnect may be recovered via `Last-Event-ID`. + +use std::{sync::Arc, time::Duration}; + +use rmcp::{ + ErrorData as McpError, RoleServer, ServerHandler, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ServerCapabilities, + ServerConfig, + }, + service::RequestContext, + transport::streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, + }, +}; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +#[derive(Clone)] +struct CancelProbe { + started: Arc, + cancelled: Arc, +} + +impl ServerHandler for CancelProbe { + #[allow(deprecated)] + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn call_tool( + &self, + _request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + self.started.notify_one(); + // Wait until the per-request cancellation token fires, or give up after a + // generous timeout so a buggy build fails via the outer assertion rather + // than hanging the test. + tokio::select! { + _ = context.ct.cancelled() => { + self.cancelled.notify_one(); + Ok(CallToolResult::success(vec![ContentBlock::text("cancelled")]).into()) + } + _ = tokio::time::sleep(Duration::from_secs(30)) => { + Ok(CallToolResult::success(vec![ContentBlock::text("ran_to_completion")]).into()) + } + } + } +} + +const CALL_BODY: &str = r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"wait_for_cancel","arguments":{}}}"#; + +struct TestServer { + url: String, + server_ct: CancellationToken, + started: Arc, + cancelled: Arc, +} + +async fn spawn_stateless_server(json_response: bool) -> anyhow::Result { + let started = Arc::new(Notify::new()); + let cancelled = Arc::new(Notify::new()); + let probe = CancelProbe { + started: started.clone(), + cancelled: cancelled.clone(), + }; + + let server_ct = CancellationToken::new(); + let config = StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_json_response(json_response) + // A short keep-alive lets the SSE server notice a dropped connection + // quickly (hyper only observes the disconnect on its next write). + .with_sse_keep_alive(Some(Duration::from_millis(100))) + .with_cancellation_token(server_ct.child_token()); + + let service: StreamableHttpService = + StreamableHttpService::new( + move || Ok(probe.clone()), + Arc::new(LocalSessionManager::default()), + config, + ); + let router = axum::Router::new().nest_service("/mcp", service); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let addr = listener.local_addr()?; + tokio::spawn({ + let ct = server_ct.clone(); + async move { + let _ = axum::serve(listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + + Ok(TestServer { + url: format!("http://{addr}/mcp"), + server_ct, + started, + cancelled, + }) +} + +/// SSE mode: the response is a stream; dropping it (client disconnect) must fire +/// the handler's cancellation token. +#[tokio::test] +async fn stateless_sse_client_disconnect_cancels_request() -> anyhow::Result<()> { + let server = spawn_stateless_server(false).await?; + let client = reqwest::Client::builder() + .pool_max_idle_per_host(0) + .build()?; + + // A single self-contained tools/call (no session, no initialize handshake). + let call = client + .post(&server.url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2025-03-26") + .body(CALL_BODY) + .send() + .await?; + assert!( + call.status().is_success(), + "tools/call failed: {:?}", + call.status() + ); + + tokio::time::timeout(Duration::from_secs(5), server.started.notified()) + .await + .expect("tool handler should start"); + + // Client disconnects mid-call: drop the streaming response (and the client). + drop(call); + drop(client); + + tokio::time::timeout(Duration::from_secs(10), server.cancelled.notified()) + .await + .expect("RequestContext::ct should fire after client disconnect (SSE)"); + + server.server_ct.cancel(); + Ok(()) +} + +/// JSON-direct mode: the server holds the connection open awaiting the single +/// response. A client that disconnects while the handler is running must still +/// fire the handler's cancellation token. +#[tokio::test] +async fn stateless_json_client_disconnect_cancels_request() -> anyhow::Result<()> { + let server = spawn_stateless_server(true).await?; + let client = reqwest::Client::builder() + .pool_max_idle_per_host(0) + .build()?; + + // In JSON mode the server does not respond until the handler completes, so + // the request stays pending; drive it from a task we can abort to disconnect. + let url = server.url.clone(); + let req_task = tokio::spawn(async move { + let _ = client + .post(&url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2025-03-26") + .body(CALL_BODY) + .send() + .await; + // Keep the client alive until the request future is dropped by abort(). + drop(client); + }); + + tokio::time::timeout(Duration::from_secs(5), server.started.notified()) + .await + .expect("tool handler should start"); + + // Client disconnects: abort the in-flight request, closing the connection. + req_task.abort(); + + tokio::time::timeout(Duration::from_secs(10), server.cancelled.notified()) + .await + .expect("RequestContext::ct should fire after client disconnect (JSON)"); + + server.server_ct.cancel(); + Ok(()) +} diff --git a/crates/rmcp/tests/test_streamable_http_event_store.rs b/crates/rmcp/tests/test_streamable_http_event_store.rs new file mode 100644 index 0000000..979ce7e --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_event_store.rs @@ -0,0 +1,372 @@ +#![cfg(all( + feature = "client", + feature = "server", + feature = "transport-streamable-http-client-reqwest", + feature = "transport-streamable-http-server", + not(feature = "local") +))] + +use std::{ + collections::HashMap, + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, +}; + +use futures::StreamExt; +use rmcp::{ + ErrorData, ServerHandler, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, + ProgressNotificationParam, ServerCapabilities, ServerConfig, + }, + service::RequestContext, + transport::streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, + session::{ + EventStore, EventStoreError, EventStream, ServerSseMessage, SessionState, SessionStore, + SessionStoreError, local::LocalSessionManager, never::NeverSessionManager, + }, + }, +}; +use tokio::sync::RwLock; +use tokio_util::sync::CancellationToken; + +#[derive(Clone, Default)] +struct InMemorySessionStore(Arc>>); + +#[async_trait::async_trait] +impl SessionStore for InMemorySessionStore { + async fn load(&self, session_id: &str) -> Result, SessionStoreError> { + Ok(self.0.read().await.get(session_id).cloned()) + } + + async fn store(&self, session_id: &str, state: &SessionState) -> Result<(), SessionStoreError> { + self.0 + .write() + .await + .insert(session_id.to_owned(), state.clone()); + Ok(()) + } + + async fn delete(&self, session_id: &str) -> Result<(), SessionStoreError> { + self.0.write().await.remove(session_id); + Ok(()) + } +} + +#[derive(Clone)] +struct StoredEvent { + stream_id: String, + event: ServerSseMessage, +} + +#[derive(Clone, Default)] +struct InMemoryEventStore { + events: Arc>>, + next_id: Arc, +} + +#[async_trait::async_trait] +impl EventStore for InMemoryEventStore { + async fn store_event( + &self, + stream_id: &str, + event: &ServerSseMessage, + ) -> Result { + let event_id = format!("event-{}", self.next_id.fetch_add(1, Ordering::Relaxed)); + let mut event = event.clone(); + event.event_id = Some(event_id.clone()); + self.events.write().await.push(StoredEvent { + stream_id: stream_id.to_owned(), + event, + }); + Ok(event_id) + } + + async fn replay_events_after( + &self, + last_event_id: &str, + ) -> Result { + let events = self.events.read().await; + let last_index = events + .iter() + .position(|stored| stored.event.event_id.as_deref() == Some(last_event_id)) + .ok_or_else(|| std::io::Error::other("event not found"))?; + let stream_id = events[last_index].stream_id.clone(); + let replay = events + .iter() + .skip(last_index + 1) + .filter(|stored| stored.stream_id == stream_id) + .map(|stored| stored.event.clone()) + .collect::>(); + Ok(Box::pin(futures::stream::iter(replay))) + } +} + +#[derive(Clone)] +struct ProgressServer; + +impl ServerHandler for ProgressServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + if request.name == "progress" || request.name == "slow-progress" { + let progress_token = context + .meta + .get_progress_token() + .expect("request includes progressToken"); + context + .peer + .notify_progress( + ProgressNotificationParam::new(progress_token, 50.0) + .with_total(100.0) + .with_message("working"), + ) + .await + .expect("progress notification is delivered"); + } + if request.name == "slow-progress" { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + Ok(CallToolResult::success(vec![ContentBlock::text("done")]).into()) + } +} + +async fn spawn_server( + session_store: Arc, + event_store: Arc, + cancellation_token: &CancellationToken, + legacy_session_mode: bool, +) -> anyhow::Result<(String, tokio::task::JoinHandle<()>)> { + let config = { + let mut config = StreamableHttpServerConfig::default(); + config.sse_keep_alive = None; + config.legacy_session_mode = legacy_session_mode; + config.cancellation_token = cancellation_token.child_token(); + config.session_store = Some(session_store); + config + }; + let router = if legacy_session_mode { + let session_manager = + Arc::new(LocalSessionManager::default().with_event_store(event_store)); + let service = StreamableHttpService::new(|| Ok(ProgressServer), session_manager, config); + axum::Router::new().nest_service("/mcp", service) + } else { + let session_manager = + Arc::new(NeverSessionManager::default().with_event_store(event_store)); + let service = StreamableHttpService::new(|| Ok(ProgressServer), session_manager, config); + axum::Router::new().nest_service("/mcp", service) + }; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let address = listener.local_addr()?; + let handle = tokio::spawn({ + let cancellation_token = cancellation_token.clone(); + async move { + let _ = axum::serve(listener, router) + .with_graceful_shutdown(async move { cancellation_token.cancelled_owned().await }) + .await; + } + }); + Ok((format!("http://{address}/mcp"), handle)) +} + +fn event_id_containing<'a>(body: &'a str, needle: &str) -> Option<&'a str> { + body.split("\n\n") + .find(|event| event.contains(needle))? + .lines() + .find_map(|line| line.strip_prefix("id: ")) +} + +#[tokio::test] +async fn restored_instance_replays_events_from_shared_store() -> anyhow::Result<()> { + let session_store: Arc = Arc::new(InMemorySessionStore::default()); + let event_store: Arc = Arc::new(InMemoryEventStore::default()); + let http = reqwest::Client::new(); + + let cancellation_a = CancellationToken::new(); + let (url_a, server_a) = spawn_server( + session_store.clone(), + event_store.clone(), + &cancellation_a, + true, + ) + .await?; + let initialize = http + .post(&url_a) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .body(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#) + .send() + .await?; + let session_id = initialize + .headers() + .get("mcp-session-id") + .expect("initialize returns a session ID") + .to_str()? + .to_owned(); + let _ = initialize.text().await?; + + let initialized_status = http + .post(&url_a) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("Mcp-Session-Id", &session_id) + .header("Mcp-Protocol-Version", "2025-06-18") + .body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#) + .send() + .await? + .status(); + assert_eq!(initialized_status, reqwest::StatusCode::ACCEPTED); + + let original_body = http + .post(&url_a) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("Mcp-Session-Id", &session_id) + .header("Mcp-Protocol-Version", "2025-06-18") + .body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"progress","arguments":{},"_meta":{"progressToken":"progress-1"}}}"#) + .send() + .await? + .text() + .await?; + let progress_event_id = event_id_containing(&original_body, "notifications/progress") + .expect("progress event has a persisted event ID") + .to_owned(); + + cancellation_a.cancel(); + server_a.await?; + + let cancellation_b = CancellationToken::new(); + let (url_b, server_b) = spawn_server(session_store, event_store, &cancellation_b, true).await?; + let replay = http + .get(&url_b) + .header("Accept", "text/event-stream") + .header("Mcp-Session-Id", &session_id) + .header("Mcp-Protocol-Version", "2025-06-18") + .header("Last-Event-ID", progress_event_id) + .send() + .await?; + assert_eq!(replay.status(), reqwest::StatusCode::OK); + let replay_body = replay.text().await?; + assert!( + replay_body.contains(r#""id":2"#), + "instance B should replay the final response stored by instance A: {replay_body}" + ); + + cancellation_b.cancel(); + server_b.await?; + Ok(()) +} + +#[tokio::test] +async fn stateless_instance_replays_events_from_shared_store() -> anyhow::Result<()> { + let session_store: Arc = Arc::new(InMemorySessionStore::default()); + let event_store = Arc::new(InMemoryEventStore::default()); + let http = reqwest::Client::new(); + + let cancellation_a = CancellationToken::new(); + let (url_a, server_a) = spawn_server( + session_store.clone(), + event_store.clone(), + &cancellation_a, + false, + ) + .await?; + let original = http + .post(&url_a) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2026-07-28") + .header("Mcp-Method", "tools/call") + .header("Mcp-Name", "slow-progress") + .body( + r#"{ + "jsonrpc":"2.0", + "id":2, + "method":"tools/call", + "params":{ + "name":"slow-progress", + "arguments":{}, + "_meta":{ + "progressToken":"progress-1", + "io.modelcontextprotocol/protocolVersion":"2026-07-28", + "io.modelcontextprotocol/clientInfo":{"name":"test","version":"1.0"}, + "io.modelcontextprotocol/clientCapabilities":{} + } + } + }"#, + ) + .send() + .await?; + assert_eq!(original.status(), reqwest::StatusCode::OK); + assert!( + original.headers().get("mcp-session-id").is_none(), + "stateless response must not create a session" + ); + let mut body = original.bytes_stream(); + let mut received = String::new(); + let progress_event_id = tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + let chunk = body + .next() + .await + .expect("response remains open until progress arrives")?; + received.push_str(&String::from_utf8_lossy(&chunk)); + if let Some(event_id) = event_id_containing(&received, "notifications/progress") { + return Ok::<_, reqwest::Error>(event_id.to_owned()); + } + } + }) + .await??; + drop(body); + + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + let stored_response = event_store.events.read().await.iter().any(|stored| { + stored.event.message.as_ref().is_some_and(|message| { + matches!( + message.as_ref(), + rmcp::model::ServerJsonRpcMessage::Response(_) + ) + }) + }); + if stored_response { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + }) + .await?; + + cancellation_a.cancel(); + server_a.await?; + + let cancellation_b = CancellationToken::new(); + let (url_b, server_b) = + spawn_server(session_store, event_store, &cancellation_b, false).await?; + let replay = http + .get(&url_b) + .header("Accept", "text/event-stream") + .header("MCP-Protocol-Version", "2026-07-28") + .header("Last-Event-ID", progress_event_id) + .send() + .await?; + assert_eq!(replay.status(), reqwest::StatusCode::OK); + let replay_body = replay.text().await?; + assert!( + replay_body.contains(r#""id":2"#), + "instance B should replay the stateless response stored by instance A: {replay_body}" + ); + + cancellation_b.cancel(); + server_b.await?; + Ok(()) +} diff --git a/crates/rmcp/tests/test_streamable_http_get_stream_auth_challenge.rs b/crates/rmcp/tests/test_streamable_http_get_stream_auth_challenge.rs new file mode 100644 index 0000000..dcd1b8d --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_get_stream_auth_challenge.rs @@ -0,0 +1,105 @@ +#![cfg(all( + feature = "transport-streamable-http-client", + feature = "transport-streamable-http-client-reqwest", + not(feature = "local") +))] + +use std::{collections::HashMap, sync::Arc}; + +use rmcp::transport::streamable_http_client::{StreamableHttpClient, StreamableHttpError}; + +/// Spin up a minimal axum server whose GET handler always responds with the given +/// status and optional `WWW-Authenticate` header — no MCP logic involved. +async fn spawn_mock_server(status: u16, www_authenticate: Option<&'static str>) -> String { + use axum::{Router, body::Body, http::Response, routing::get}; + + let router = Router::new().route( + "/mcp", + get(move || async move { + let mut builder = Response::builder().status(status); + if let Some(challenge) = www_authenticate { + builder = builder.header("www-authenticate", challenge); + } + builder.body(Body::empty()).unwrap() + }), + ); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + }); + + format!("http://{addr}/mcp") +} + +async fn get_stream_against( + status: u16, + www_authenticate: Option<&'static str>, +) -> Result<(), StreamableHttpError> { + let url = spawn_mock_server(status, www_authenticate).await; + reqwest::Client::new() + .get_stream(Arc::from(url.as_str()), None, None, None, HashMap::new()) + .await + .map(|_| ()) +} + +/// A 401 carrying a `WWW-Authenticate` challenge must surface as `AuthRequired`, +/// which is the variant `AuthClient::call_reacting_to_challenges` catches to +/// refresh the token and retry. Classified as a plain `Client` error instead, the +/// refresh never runs and the stream just fails. +#[tokio::test] +async fn get_stream_maps_401_challenge_to_auth_required() { + let result = get_stream_against(401, Some("Bearer realm=\"mcp\"")).await; + + match result { + Err(StreamableHttpError::AuthRequired(err)) => { + assert_eq!(err.www_authenticate_header, "Bearer realm=\"mcp\""); + } + other => panic!("expected AuthRequired, got: {other:?}"), + } +} + +/// A 403 carrying a challenge must surface as `InsufficientScope`, with the +/// required scope extracted from the header. +#[tokio::test] +async fn get_stream_maps_403_challenge_to_insufficient_scope() { + let result = get_stream_against( + 403, + Some("Bearer error=\"insufficient_scope\", scope=\"mcp:read\""), + ) + .await; + + match result { + Err(StreamableHttpError::InsufficientScope(err)) => { + assert_eq!(err.required_scope.as_deref(), Some("mcp:read")); + } + other => panic!("expected InsufficientScope, got: {other:?}"), + } +} + +/// Without a `WWW-Authenticate` header there is no challenge to act on, so a 401 +/// must keep falling through to the ordinary error path rather than being +/// reported as an auth challenge the caller can retry. +#[tokio::test] +async fn get_stream_401_without_challenge_is_not_auth_required() { + let result = get_stream_against(401, None).await; + + assert!( + !matches!(result, Err(StreamableHttpError::AuthRequired(_))), + "401 with no challenge header must not be classified as AuthRequired" + ); + assert!(result.is_err(), "401 must still be an error"); +} + +/// 405 keeps its dedicated meaning — the server does not support SSE on GET — +/// and must not be swallowed by the new auth branches. +#[tokio::test] +async fn get_stream_405_still_reports_sse_unsupported() { + let result = get_stream_against(405, None).await; + + assert!( + matches!(result, Err(StreamableHttpError::ServerDoesNotSupportSse)), + "expected ServerDoesNotSupportSse, got: {result:?}" + ); +} diff --git a/crates/rmcp/tests/test_streamable_http_json_response.rs b/crates/rmcp/tests/test_streamable_http_json_response.rs index 09dd69c..62daa7a 100644 --- a/crates/rmcp/tests/test_streamable_http_json_response.rs +++ b/crates/rmcp/tests/test_streamable_http_json_response.rs @@ -1,6 +1,14 @@ #![cfg(not(feature = "local"))] -use rmcp::transport::streamable_http_server::{ - StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, +use rmcp::{ + ErrorData, ServerHandler, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, + ProgressNotificationParam, ServerCapabilities, ServerConfig, + }, + service::RequestContext, + transport::streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, + }, }; use tokio_util::sync::CancellationToken; @@ -8,6 +16,74 @@ mod common; use common::calculator::Calculator; const INIT_BODY: &str = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#; +const CALL_WITH_PROGRESS_BODY: &str = r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"progress","arguments":{},"_meta":{"progressToken":"progress-test-1"}}}"#; +const NEGOTIATED_CALL_BODY: &str = r#"{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": "terminal", + "arguments": {}, + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": { + "name": "test", + "version": "1.0" + }, + "io.modelcontextprotocol/clientCapabilities": {} + } + } +}"#; +const NEGOTIATED_CALL_WITH_PROGRESS_BODY: &str = r#"{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": "progress", + "arguments": {}, + "_meta": { + "progressToken": "progress-test-1", + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": { + "name": "test", + "version": "1.0" + }, + "io.modelcontextprotocol/clientCapabilities": {} + } + } +}"#; + +#[derive(Clone)] +struct ProgressServer; + +impl ServerHandler for ProgressServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + if request.name == "progress" { + let progress_token = context + .meta + .get_progress_token() + .expect("request includes progressToken"); + context + .peer + .notify_progress( + ProgressNotificationParam::new(progress_token, 50.0) + .with_total(100.0) + .with_message("working"), + ) + .await + .expect("progress notification is delivered"); + } + Ok(CallToolResult::success(vec![ContentBlock::text("done")]).into()) + } +} async fn spawn_server( config: StreamableHttpServerConfig, @@ -34,12 +110,37 @@ async fn spawn_server( (client, base_url, ct) } +async fn spawn_progress_server( + config: StreamableHttpServerConfig, +) -> (reqwest::Client, String, CancellationToken) { + let ct = config.cancellation_token.clone(); + let service: StreamableHttpService = + StreamableHttpService::new(|| Ok(ProgressServer), Default::default(), config); + + let router = axum::Router::new().nest_service("/mcp", service); + let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = tcp_listener.local_addr().unwrap(); + + tokio::spawn({ + let ct = ct.clone(); + async move { + let _ = axum::serve(tcp_listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + + let client = reqwest::Client::new(); + let base_url = format!("http://{addr}/mcp"); + (client, base_url, ct) +} + #[tokio::test] async fn stateless_json_response_returns_application_json() -> anyhow::Result<()> { let ct = CancellationToken::new(); let (client, url, ct) = spawn_server( StreamableHttpServerConfig::default() - .with_stateful_mode(false) + .with_legacy_session_mode(false) .with_json_response(true) .with_sse_keep_alive(None) .with_cancellation_token(ct.child_token()), @@ -76,12 +177,162 @@ async fn stateless_json_response_returns_application_json() -> anyhow::Result<() Ok(()) } +#[tokio::test] +async fn stateless_negotiated_terminal_response_returns_application_json() -> anyhow::Result<()> { + let ct = CancellationToken::new(); + let (client, url, ct) = spawn_progress_server( + StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_json_response(true) + .with_sse_keep_alive(None) + .with_cancellation_token(ct.child_token()), + ) + .await; + + let response = client + .post(&url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2026-07-28") + .header("Mcp-Method", "tools/call") + .header("Mcp-Name", "terminal") + .body(NEGOTIATED_CALL_BODY) + .send() + .await?; + + assert_eq!(response.status(), 200); + + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or(""); + assert!( + content_type.contains("application/json"), + "Expected application/json, got: {content_type}" + ); + + let body: serde_json::Value = response.json().await?; + assert_eq!(body["id"], 2); + assert!(body["result"].is_object(), "Expected result object"); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn stateless_json_response_falls_back_to_sse_for_progress() -> anyhow::Result<()> { + let ct = CancellationToken::new(); + let (client, url, ct) = spawn_progress_server( + StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_json_response(true) + .with_sse_keep_alive(None) + .with_cancellation_token(ct.child_token()), + ) + .await; + + let response = client + .post(&url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .body(CALL_WITH_PROGRESS_BODY) + .send() + .await?; + + assert_eq!(response.status(), 200); + + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or(""); + assert!( + content_type.contains("text/event-stream"), + "Expected SSE fallback, got: {content_type}" + ); + + let body = response.text().await?; + let messages: Vec = body + .lines() + .filter_map(|line| line.strip_prefix("data:")) + .map(str::trim) + .filter(|data| !data.is_empty()) + .map(serde_json::from_str) + .collect::>()?; + assert_eq!(messages.len(), 2, "Expected progress and result: {body}"); + assert_eq!(messages[0]["method"], "notifications/progress"); + assert_eq!(messages[1]["id"], 2); + assert!( + messages[1]["result"].is_object(), + "Expected result object: {body}" + ); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn stateless_negotiated_json_response_falls_back_to_sse_for_progress() -> anyhow::Result<()> { + let ct = CancellationToken::new(); + let (client, url, ct) = spawn_progress_server( + StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_json_response(true) + .with_sse_keep_alive(None) + .with_cancellation_token(ct.child_token()), + ) + .await; + + let response = client + .post(&url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2026-07-28") + .header("Mcp-Method", "tools/call") + .header("Mcp-Name", "progress") + .body(NEGOTIATED_CALL_WITH_PROGRESS_BODY) + .send() + .await?; + + assert_eq!(response.status(), 200); + + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or(""); + assert!( + content_type.contains("text/event-stream"), + "Expected SSE fallback, got: {content_type}" + ); + + let body = response.text().await?; + let messages: Vec = body + .lines() + .filter_map(|line| line.strip_prefix("data:")) + .map(str::trim) + .filter(|data| !data.is_empty()) + .map(serde_json::from_str) + .collect::>()?; + assert_eq!(messages.len(), 2, "Expected progress and result: {body}"); + assert_eq!(messages[0]["method"], "notifications/progress"); + assert_eq!(messages[1]["id"], 2); + assert!( + messages[1]["result"].is_object(), + "Expected result object: {body}" + ); + + ct.cancel(); + Ok(()) +} + #[tokio::test] async fn stateless_sse_mode_default_unchanged() -> anyhow::Result<()> { let ct = CancellationToken::new(); let (client, url, ct) = spawn_server( StreamableHttpServerConfig::default() - .with_stateful_mode(false) + .with_legacy_session_mode(false) .with_sse_keep_alive(None) .with_cancellation_token(ct.child_token()), ) @@ -118,9 +369,9 @@ async fn stateless_sse_mode_default_unchanged() -> anyhow::Result<()> { } #[tokio::test] -async fn json_response_ignored_in_stateful_mode() -> anyhow::Result<()> { +async fn json_response_ignored_in_legacy_session_mode() -> anyhow::Result<()> { let ct = CancellationToken::new(); - // json_response: true has no effect when stateful_mode: true — server still uses SSE + // json_response: true has no effect when legacy_session_mode: true — server still uses SSE let (client, url, ct) = spawn_server( StreamableHttpServerConfig::default() .with_json_response(true) diff --git a/crates/rmcp/tests/test_streamable_http_malformed_json_request.rs b/crates/rmcp/tests/test_streamable_http_malformed_json_request.rs new file mode 100644 index 0000000..c34e96e --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_malformed_json_request.rs @@ -0,0 +1,195 @@ +#![cfg(all( + feature = "transport-streamable-http-client", + feature = "transport-streamable-http-client-reqwest", + feature = "client", + not(feature = "local") +))] + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use axum::{Router, body::Bytes, http::StatusCode, response::IntoResponse, routing::post}; +use rmcp::{ + ServiceError, ServiceExt, + model::{ + CallToolRequestParams, ClientJsonRpcMessage, ClientNotification, ClientRequest, + InitializedNotification, PingRequest, RequestId, + }, + transport::{ + StreamableHttpClientTransport, + streamable_http_client::{ + StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse, + }, + }, +}; +use rstest::rstest; + +/// Initialize succeeds; every later POST (including `tools/call`) returns +/// HTTP 200 + `application/json` with a body that is not a JSON-RPC message. +async fn spawn_malformed_json_request_server(call_body: &'static str) -> String { + let router = Router::new().route( + "/mcp", + post(move |body: Bytes| async move { + let parsed: serde_json::Value = serde_json::from_slice(&body).unwrap_or_default(); + let method = parsed.get("method").and_then(|m| m.as_str()).unwrap_or(""); + if method == "initialize" { + let id = parsed.get("id").cloned().unwrap_or(serde_json::json!(1)); + return ( + StatusCode::OK, + [ + (http::header::CONTENT_TYPE, "application/json"), + ( + http::HeaderName::from_static("mcp-session-id"), + "test-session", + ), + ], + serde_json::json!({ + "jsonrpc": "2.0", + "id": id, + "result": { + "protocolVersion": "2025-03-26", + "capabilities": { "tools": {} }, + "serverInfo": { "name": "mock", "version": "0.0.1" } + } + }) + .to_string(), + ) + .into_response(); + } + if method == "notifications/initialized" { + return StatusCode::ACCEPTED.into_response(); + } + ( + StatusCode::OK, + [(http::header::CONTENT_TYPE, "application/json")], + call_body, + ) + .into_response() + }), + ); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + }); + format!("http://{addr}/mcp") +} + +async fn spawn_json_body_server(body: &'static str) -> String { + let router = Router::new().route( + "/mcp", + post(move || async move { + ( + StatusCode::OK, + [(http::header::CONTENT_TYPE, "application/json")], + body, + ) + }), + ); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + }); + format!("http://{addr}/mcp") +} + +/// Public-API repro: `call_tool` must return a transport error instead of +/// hanging when the server answers a request POST with 200 JSON that is not +/// a JSON-RPC message. +#[rstest] +#[case::empty_object("{}")] +#[case::empty_body("")] +#[case::invalid_json("{")] +#[case::html("not json")] +#[tokio::test] +async fn call_tool_errors_on_malformed_json_200(#[case] call_body: &'static str) { + let url = spawn_malformed_json_request_server(call_body).await; + let transport = StreamableHttpClientTransport::from_uri(url); + let client = ().serve(transport).await.expect("initialize should succeed"); + let peer = client.peer().clone(); + + let result = tokio::time::timeout( + Duration::from_secs(2), + peer.call_tool(CallToolRequestParams::new("anything")), + ) + .await + .expect("call_tool must return instead of hanging on malformed JSON 200"); + + match result { + Err(ServiceError::TransportSend(ref dyn_err)) => { + let err_msg = format!("{dyn_err}"); + assert!( + err_msg.contains("unexpected server response"), + "expected UnexpectedServerResponse, got: {err_msg}" + ); + let preview = if call_body.is_empty() { + "" + } else { + call_body + }; + assert!( + err_msg.contains(preview), + "expected body preview {preview:?} in error, got: {err_msg}" + ); + } + other => panic!("expected TransportSend(UnexpectedServerResponse), got: {other:?}"), + } + + let _ = client.cancel().await; +} + +/// Direct `post_message` check: a request POST cannot treat a JSON parse +/// failure as Accepted. +#[tokio::test] +async fn post_request_malformed_json_is_unexpected_server_response() { + let url = spawn_json_body_server("{}").await; + let client = reqwest::Client::new(); + let result = client + .post_message( + Arc::from(url.as_str()), + ClientJsonRpcMessage::request( + ClientRequest::PingRequest(PingRequest::default()), + RequestId::Number(1), + ), + None, + None, + HashMap::new(), + ) + .await; + + match result { + Err(StreamableHttpError::UnexpectedServerResponse(ref msg)) => { + assert!( + msg.contains("{}"), + "expected body preview in error, got: {msg}" + ); + } + other => panic!("expected UnexpectedServerResponse, got: {other:?}"), + } +} + +/// Notification POSTs still treat an unusable JSON body as Accepted. +/// That fallback is for messages that do not await a reply. +#[tokio::test] +async fn post_notification_malformed_json_is_still_accepted() { + let url = spawn_json_body_server("{}").await; + let client = reqwest::Client::new(); + let result = client + .post_message( + Arc::from(url.as_str()), + ClientJsonRpcMessage::notification(ClientNotification::InitializedNotification( + InitializedNotification::default(), + )), + None, + None, + HashMap::new(), + ) + .await; + + match result { + Ok(StreamableHttpPostResponse::Accepted) => {} + other => panic!("expected Accepted, got: {other:?}"), + } +} diff --git a/crates/rmcp/tests/test_streamable_http_priming.rs b/crates/rmcp/tests/test_streamable_http_priming.rs index 436d482..1e0e5e5 100644 --- a/crates/rmcp/tests/test_streamable_http_priming.rs +++ b/crates/rmcp/tests/test_streamable_http_priming.rs @@ -14,7 +14,7 @@ use common::calculator::Calculator; async fn test_priming_on_stream_start() -> anyhow::Result<()> { let ct = CancellationToken::new(); - // stateful_mode: true automatically enables priming with DEFAULT_RETRY_INTERVAL (3 seconds) + // legacy_session_mode: true automatically enables priming with DEFAULT_RETRY_INTERVAL (3 seconds) let service: StreamableHttpService = StreamableHttpService::new( || Ok(Calculator::new()), @@ -416,7 +416,7 @@ async fn test_priming_on_stream_close() -> anyhow::Result<()> { let ct = CancellationToken::new(); let session_manager = Arc::new(LocalSessionManager::default()); - // stateful_mode: true automatically enables priming with DEFAULT_RETRY_INTERVAL (3 seconds) + // legacy_session_mode: true automatically enables priming with DEFAULT_RETRY_INTERVAL (3 seconds) let service = StreamableHttpService::new( || Ok(Calculator::new()), session_manager.clone(), diff --git a/crates/rmcp/tests/test_streamable_http_protocol_version.rs b/crates/rmcp/tests/test_streamable_http_protocol_version.rs index 0ed61c0..76005a7 100644 --- a/crates/rmcp/tests/test_streamable_http_protocol_version.rs +++ b/crates/rmcp/tests/test_streamable_http_protocol_version.rs @@ -1,10 +1,17 @@ #![cfg(not(feature = "local"))] -//! Regression tests for the `MCP-Protocol-Version` header / initialize body consistency check. -use std::sync::Arc; - -use rmcp::transport::streamable_http_server::{ - StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, +//! Streamable HTTP protocol-version and request-metadata validation tests. +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, }; + +use rmcp::{ + ServerHandler, + transport::streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, + }, +}; +use serde_json::{Value, json}; use tokio_util::sync::CancellationToken; mod common; @@ -50,7 +57,7 @@ async fn spawn_server_with_manager( fn stateless_json_config() -> StreamableHttpServerConfig { StreamableHttpServerConfig::default() - .with_stateful_mode(false) + .with_legacy_session_mode(false) .with_json_response(true) .with_sse_keep_alive(None) .with_cancellation_token(CancellationToken::new()) @@ -58,7 +65,7 @@ fn stateless_json_config() -> StreamableHttpServerConfig { fn stateful_config() -> StreamableHttpServerConfig { StreamableHttpServerConfig::default() - .with_stateful_mode(true) + .with_legacy_session_mode(true) .with_sse_keep_alive(None) .with_cancellation_token(CancellationToken::new()) } @@ -80,6 +87,17 @@ async fn post_init( req.send().await.expect("send initialize request") } +/// First JSON-RPC message of an SSE response, skipping the empty priming event. +async fn sse_payload(response: reqwest::Response) -> anyhow::Result { + let body = response.text().await?; + let data = body + .lines() + .filter_map(|line| line.strip_prefix("data: ")) + .find(|data| !data.is_empty()) + .expect("response must carry a JSON-RPC payload"); + Ok(serde_json::from_str(data)?) +} + async fn post_non_initialize(client: &reqwest::Client, url: &str) -> reqwest::Response { client .post(url) @@ -91,6 +109,31 @@ async fn post_non_initialize(client: &reqwest::Client, url: &str) -> reqwest::Re .expect("send non-initialize request") } +async fn post_modern_request( + client: &reqwest::Client, + url: &str, + method: &str, + name: Option<&str>, + params: Value, +) -> reqwest::Response { + let mut request = client + .post(url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2026-07-28") + .header("Mcp-Method", method) + .json(&json!({ + "jsonrpc": "2.0", + "id": 1, + "method": method, + "params": params, + })); + if let Some(name) = name { + request = request.header("Mcp-Name", name); + } + request.send().await.expect("send modern request") +} + #[tokio::test] async fn stateless_init_rejects_when_header_older_than_body() -> anyhow::Result<()> { let (client, url, ct) = spawn_server(stateless_json_config()).await; @@ -154,6 +197,27 @@ async fn stateless_init_accepts_when_header_absent() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn stateless_init_naming_a_post_handshake_version_negotiates_down() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server(stateless_json_config()).await; + + let response = post_init(&client, &url, None, "2026-07-28").await; + assert_eq!( + response.status(), + 200, + "initialize selects legacy semantics whatever version it names" + ); + + let body: serde_json::Value = response.json().await?; + assert_eq!( + body["result"]["protocolVersion"], "2025-11-25", + "initialize must settle on a version that still has the handshake" + ); + + ct.cancel(); + Ok(()) +} + #[tokio::test] async fn stateful_init_rejects_when_header_mismatches_body() -> anyhow::Result<()> { let (client, url, ct) = spawn_server(stateful_config()).await; @@ -185,3 +249,668 @@ async fn stateful_rejected_initial_posts_do_not_create_sessions() -> anyhow::Res ct.cancel(); Ok(()) } + +/// The `initialize` handshake exists only in the revisions before `2026-07-28`, +/// so naming a later version in it does not make the request a modern one: the +/// server keeps legacy semantics, opens a session, and answers with a version it +/// can actually serve over the handshake. +#[tokio::test] +async fn stateful_init_naming_a_post_handshake_version_opens_a_legacy_session() -> anyhow::Result<()> +{ + let (client, url, ct) = spawn_server(stateful_config()).await; + + let response = post_init(&client, &url, None, "2026-07-28").await; + assert_eq!( + response.status(), + 200, + "initialize selects legacy semantics whatever version it names" + ); + assert!( + response.headers().contains_key("Mcp-Session-Id"), + "the handshake must open a session, got headers: {:?}", + response.headers() + ); + + let payload = sse_payload(response).await?; + assert_eq!( + payload["result"]["protocolVersion"], "2025-11-25", + "initialize must settle on a version that still has the handshake" + ); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn stateless_missing_protocol_header_returns_header_mismatch() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server(stateless_json_config()).await; + + let body = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"_meta":{"io.modelcontextprotocol/protocolVersion":"2026-07-28","io.modelcontextprotocol/clientInfo":{"name":"test","version":"1.0"},"io.modelcontextprotocol/clientCapabilities":{}}}}"#; + let response = client + .post(&url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("Mcp-Method", "tools/list") + .body(body) + .send() + .await + .expect("send non-initialize request"); + + assert_eq!(response.status(), 400); + + let body: serde_json::Value = response.json().await?; + assert_eq!(body["jsonrpc"], "2.0"); + assert_eq!(body["id"], 1); + assert_eq!(body["error"]["code"], -32020); + assert!( + body["error"]["message"] + .as_str() + .is_some_and(|message| message.contains("requires MCP-Protocol-Version header")), + "expected missing protocol header message, got: {body}" + ); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn stateless_tools_list_rejects_missing_request_meta() { + let (client, url, ct) = spawn_server(stateless_json_config()).await; + + let response = post_modern_request(&client, &url, "tools/list", None, json!({})).await; + + assert_eq!(response.status(), 400); + let body: Value = response.json().await.expect("response should be JSON"); + assert_eq!(body["error"]["code"], -32602); + assert!( + body["error"]["message"] + .as_str() + .is_some_and(|message| message.contains("io.modelcontextprotocol/protocolVersion")), + "expected error message to mention protocolVersion, got: {body}" + ); + ct.cancel(); +} + +#[tokio::test] +async fn stateless_tools_call_rejects_missing_request_meta() { + let (client, url, ct) = spawn_server(stateless_json_config()).await; + + let response = post_modern_request( + &client, + &url, + "tools/call", + Some("sum"), + json!({ + "name": "sum", + "arguments": { + "a": 1, + "b": 2 + } + }), + ) + .await; + + assert_eq!(response.status(), 400); + let body: Value = response.json().await.expect("response should be JSON"); + assert_eq!(body["error"]["code"], -32602); + assert!( + body["error"]["message"] + .as_str() + .is_some_and(|message| message.contains("io.modelcontextprotocol/protocolVersion")), + "expected error message to mention protocolVersion, got: {body}" + ); + ct.cancel(); +} + +#[tokio::test] +async fn stateless_request_rejects_missing_meta_protocol_version() { + let (client, url, ct) = spawn_server(stateless_json_config()).await; + + let response = post_modern_request( + &client, + &url, + "tools/list", + None, + json!({ + "_meta": { + "io.modelcontextprotocol/clientInfo": { + "name": "test", + "version": "1.0" + }, + "io.modelcontextprotocol/clientCapabilities": {} + } + }), + ) + .await; + + assert_eq!(response.status(), 400); + let body: Value = response.json().await.expect("response should be JSON"); + assert_eq!(body["error"]["code"], -32602); + assert!( + body["error"]["message"] + .as_str() + .is_some_and(|message| message.contains("io.modelcontextprotocol/protocolVersion")), + "expected error message to mention protocolVersion, got: {body}" + ); + ct.cancel(); +} + +#[tokio::test] +async fn stateless_request_rejects_missing_meta_client_capabilities() { + let (client, url, ct) = spawn_server(stateless_json_config()).await; + + let response = post_modern_request( + &client, + &url, + "tools/list", + None, + json!({ + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": { + "name": "test", + "version": "1.0" + } + } + }), + ) + .await; + + assert_eq!(response.status(), 400); + let body: Value = response.json().await.expect("response should be JSON"); + assert_eq!(body["error"]["code"], -32602); + assert!( + body["error"]["message"] + .as_str() + .is_some_and(|message| message.contains("io.modelcontextprotocol/clientCapabilities")), + "expected error message to mention clientCapabilities, got: {body}" + ); + ct.cancel(); +} + +#[tokio::test] +async fn stateless_request_accepts_missing_optional_meta_client_info() { + let (client, url, ct) = spawn_server(stateless_json_config()).await; + + let response = post_modern_request( + &client, + &url, + "tools/list", + None, + json!({ + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": {} + } + }), + ) + .await; + + assert_eq!(response.status(), 200); + ct.cancel(); +} + +// --------------------------------------------------------------------------- +// Opt-in seam: `with_stateless_protocol_metadata_required(true)` +// --------------------------------------------------------------------------- +// In stateless mode, every non-initialize Streamable HTTP JSON-RPC request +// POST must carry the `MCP-Protocol-Version` HTTP header (missing → HTTP 400 / +// JSON-RPC `-32020` HeaderMismatch before dispatch). Every non-initialize, +// non-discover request must additionally carry a per-request +// `io.modelcontextprotocol/protocolVersion` in `_meta` (missing → HTTP 400 / +// JSON-RPC `-32602`). The existing rule that requires `server/discover` to +// carry `_meta.protocolVersion` is preserved unchanged. +// +// Non-dispatch is proven with an explicit invocation counter (`AtomicUsize`). + +#[derive(Clone)] +struct CountingServer { + lists: Arc, +} + +impl CountingServer { + fn new() -> (Self, Arc) { + let lists = Arc::new(AtomicUsize::new(0)); + ( + Self { + lists: lists.clone(), + }, + lists, + ) + } +} + +impl ServerHandler for CountingServer { + fn get_info(&self) -> rmcp::model::ServerConfig { + rmcp::model::ServerConfig::new( + rmcp::model::ServerCapabilities::builder() + .enable_tools() + .build(), + ) + } + + fn list_tools( + &self, + _request: Option, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + + Send + + '_ { + self.lists.fetch_add(1, Ordering::SeqCst); + std::future::ready(Ok(rmcp::model::ListToolsResult::default())) + } +} + +async fn spawn_counting( + config: StreamableHttpServerConfig, +) -> (reqwest::Client, String, CancellationToken, Arc) { + let ct = config.cancellation_token.clone(); + let (server, lists) = CountingServer::new(); + let service: StreamableHttpService = + StreamableHttpService::new(move || Ok(server.clone()), Default::default(), config); + + let router = axum::Router::new().nest_service("/mcp", service); + let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = tcp_listener.local_addr().unwrap(); + + tokio::spawn({ + let ct = ct.clone(); + async move { + let _ = axum::serve(tcp_listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + + ( + reqwest::Client::new(), + format!("http://{addr}/mcp"), + ct, + lists, + ) +} + +fn modern_required_config() -> StreamableHttpServerConfig { + StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_json_response(true) + .with_sse_keep_alive(None) + .with_stateless_protocol_metadata_required(true) + .with_cancellation_token(CancellationToken::new()) +} + +async fn post_seam( + client: &reqwest::Client, + url: &str, + body: &str, + header_version: Option<&str>, + extra_headers: &[(&str, &str)], +) -> reqwest::Response { + let mut req = client + .post(url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .body(body.to_owned()); + if let Some(h) = header_version { + req = req.header("MCP-Protocol-Version", h); + } + for (k, v) in extra_headers { + req = req.header(*k, *v); + } + req.send().await.expect("send request") +} + +// With the seam disabled, explicit stateless mode still dispatches a request +// without protocol metadata, preserving backwards compatibility. +#[tokio::test] +async fn seam_disabled_preserves_stateless_compatibility() -> anyhow::Result<()> { + let (client, url, ct, lists) = spawn_counting(stateless_json_config()).await; + + let body = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#; + let response = post_seam(&client, &url, body, None, &[]).await; + assert_eq!(response.status(), 200); + assert_eq!( + lists.load(Ordering::SeqCst), + 1, + "seam-disabled stateless config must dispatch exactly once" + ); + + ct.cancel(); + Ok(()) +} + +// With legacy compatibility enabled, the seam applies only to requests that +// are routed statelessly: an established legacy session remains unchanged, +// while a self-identifying modern request is still enforced. +#[tokio::test] +async fn seam_mixed_mode_enforces_only_stateless_routed_requests() -> anyhow::Result<()> { + let config = StreamableHttpServerConfig::default() + .with_legacy_session_mode(true) + .with_json_response(true) + .with_sse_keep_alive(None) + .with_stateless_protocol_metadata_required(true) + .with_cancellation_token(CancellationToken::new()); + let (client, url, ct, lists) = spawn_counting(config).await; + + let legacy_body = r#"{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}"#; + let unknown_response = post_seam( + &client, + &url, + legacy_body, + None, + &[("Mcp-Session-Id", "unknown-session")], + ) + .await; + assert_eq!( + unknown_response.status(), + reqwest::StatusCode::NOT_FOUND, + "the stateless seam must not shadow the legacy unknown-session boundary" + ); + assert_eq!(lists.load(Ordering::SeqCst), 0); + + let initialize = post_init(&client, &url, None, "2025-11-25").await; + assert_eq!(initialize.status(), 200); + let session_id = initialize + .headers() + .get("Mcp-Session-Id") + .and_then(|value| value.to_str().ok()) + .expect("legacy initialize response must include a session id") + .to_owned(); + + let legacy_response = post_seam( + &client, + &url, + legacy_body, + None, + &[("Mcp-Session-Id", &session_id)], + ) + .await; + assert_eq!(legacy_response.status(), 200); + let _ = legacy_response.bytes().await?; + assert_eq!( + lists.load(Ordering::SeqCst), + 1, + "legacy session routing must remain unchanged" + ); + + let delete_response = client + .delete(&url) + .header("Mcp-Session-Id", &session_id) + .send() + .await?; + assert_eq!(delete_response.status(), reqwest::StatusCode::ACCEPTED); + + let terminated_response = post_seam( + &client, + &url, + legacy_body, + None, + &[("Mcp-Session-Id", &session_id)], + ) + .await; + assert_eq!( + terminated_response.status(), + reqwest::StatusCode::NOT_FOUND, + "the stateless seam must not shadow the legacy terminated-session boundary" + ); + assert_eq!(lists.load(Ordering::SeqCst), 1); + + let modern_body = r#"{"jsonrpc":"2.0","id":3,"method":"tools/list","params":{}}"#; + let modern_response = post_seam( + &client, + &url, + modern_body, + Some("2026-07-28"), + &[("Mcp-Method", "tools/list")], + ) + .await; + assert_eq!(modern_response.status(), 400); + let payload: serde_json::Value = modern_response.json().await?; + assert_eq!(payload["error"]["code"], -32602); + assert_eq!( + lists.load(Ordering::SeqCst), + 1, + "stateless metadata rejection must happen before another handler dispatch" + ); + + ct.cancel(); + Ok(()) +} + +// Both signals absent → missing header / -32020 before dispatch. +#[tokio::test] +async fn seam_opt_in_rejects_missing_header_before_dispatch() -> anyhow::Result<()> { + let (client, url, ct, lists) = spawn_counting(modern_required_config()).await; + + let without_meta = r#"{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}"#; + let response = post_seam(&client, &url, without_meta, None, &[]).await; + assert_eq!(response.status(), 400, "missing header must yield HTTP 400"); + let payload: serde_json::Value = response.json().await?; + assert_eq!(payload["error"]["code"], -32020); + assert_eq!( + lists.load(Ordering::SeqCst), + 0, + "handler invocation counter must stay at 0 when the header is missing" + ); + + ct.cancel(); + Ok(()) +} + +// Header present but `_meta.protocolVersion` absent → -32602 before +// dispatch. Counter must equal 0. +#[tokio::test] +async fn seam_opt_in_rejects_missing_meta_before_dispatch() -> anyhow::Result<()> { + let (client, url, ct, lists) = spawn_counting(modern_required_config()).await; + + let body = r#"{"jsonrpc":"2.0","id":4,"method":"tools/list","params":{}}"#; + let response = post_seam( + &client, + &url, + body, + Some("2026-07-28"), + &[("Mcp-Method", "tools/list")], + ) + .await; + assert_eq!(response.status(), 400); + let payload: serde_json::Value = response.json().await?; + assert_eq!(payload["error"]["code"], -32602); + assert_eq!( + lists.load(Ordering::SeqCst), + 0, + "handler invocation counter must stay at 0 when _meta.protocolVersion is missing" + ); + + ct.cancel(); + Ok(()) +} + +// A 2026 request with protocolVersion but no clientCapabilities reaches the +// existing inline-metadata validator and is rejected before method dispatch. +#[tokio::test] +async fn seam_opt_in_rejects_missing_client_capabilities_before_dispatch() -> anyhow::Result<()> { + let (client, url, ct, lists) = spawn_counting(modern_required_config()).await; + + let body = r#"{"jsonrpc":"2.0","id":5,"method":"tools/list","params":{"_meta":{"io.modelcontextprotocol/protocolVersion":"2026-07-28"}}}"#; + let response = post_seam( + &client, + &url, + body, + Some("2026-07-28"), + &[("Mcp-Method", "tools/list")], + ) + .await; + assert_eq!(response.status(), 400); + let payload: serde_json::Value = response.json().await?; + assert_eq!(payload["error"]["code"], -32602); + assert_eq!( + lists.load(Ordering::SeqCst), + 0, + "handler invocation counter must stay at 0 when clientCapabilities is missing" + ); + + ct.cancel(); + Ok(()) +} + +// Both signals present at the current version, plus the routing header +// required by `validate_standard_headers` for `>= STANDARD_HEADERS`. The +// seam must dispatch and the handler must run exactly once. +#[tokio::test] +async fn seam_opt_in_dispatches_when_header_and_meta_present() -> anyhow::Result<()> { + let (client, url, ct, lists) = spawn_counting(modern_required_config()).await; + + let body = r#"{"jsonrpc":"2.0","id":6,"method":"tools/list","params":{"_meta":{"io.modelcontextprotocol/protocolVersion":"2026-07-28","io.modelcontextprotocol/clientInfo":{"name":"test","version":"1.0"},"io.modelcontextprotocol/clientCapabilities":{}}}}"#; + let response = post_seam( + &client, + &url, + body, + Some("2026-07-28"), + &[("Mcp-Method", "tools/list")], + ) + .await; + assert_eq!(response.status(), 200); + let payload: serde_json::Value = response.json().await?; + assert!(payload.get("result").is_some()); + assert_eq!( + lists.load(Ordering::SeqCst), + 1, + "handler invocation counter must equal 1 on successful dispatch" + ); + + ct.cancel(); + Ok(()) +} + +// rmcp clients negotiated below 2026-07-28 send the protocol header but do +// not attach per-request protocol metadata. Requiring the metadata therefore +// rejects their real request shape even if the handler supports that version. +#[tokio::test] +async fn seam_opt_in_rejects_older_rmcp_client_request_shape() -> anyhow::Result<()> { + let (client, url, ct, lists) = spawn_counting(modern_required_config()).await; + + let body = r#"{"jsonrpc":"2.0","id":7,"method":"tools/list","params":{}}"#; + let response = post_seam(&client, &url, body, Some("2025-11-25"), &[]).await; + assert_eq!(response.status(), 400); + let payload: serde_json::Value = response.json().await?; + assert_eq!(payload["error"]["code"], -32602); + assert_eq!( + lists.load(Ordering::SeqCst), + 0, + "an older client request without per-request metadata must not dispatch" + ); + + ct.cancel(); + Ok(()) +} + +// When the 2026+ header is present but `Mcp-Method` is missing, the +// existing `validate_standard_headers` rule must fire first with -32020, +// before the seam's -32602 missing-meta check. +#[tokio::test] +async fn seam_opt_in_preserves_standard_header_precedence() -> anyhow::Result<()> { + let (client, url, ct, lists) = spawn_counting(modern_required_config()).await; + + let body = r#"{"jsonrpc":"2.0","id":7,"method":"tools/list","params":{"_meta":{"io.modelcontextprotocol/protocolVersion":"2026-07-28"}}}"#; + let response = post_seam(&client, &url, body, Some("2026-07-28"), &[]).await; + assert_eq!(response.status(), 400); + let payload: serde_json::Value = response.json().await?; + assert_eq!( + payload["error"]["code"], -32020, + "standard-headers routing-header error must precede the seam's meta check" + ); + assert_eq!( + lists.load(Ordering::SeqCst), + 0, + "handler invocation counter must stay at 0 when a routing header is missing" + ); + + ct.cancel(); + Ok(()) +} + +// `initialize` is exempt from the new required-header check while retaining +// its existing optional-header and header/body consistency rules. +#[tokio::test] +async fn seam_opt_in_preserves_initialize_rules() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server(modern_required_config()).await; + + let response = post_init(&client, &url, None, "2025-11-25").await; + assert_eq!(response.status(), 200); + + let response = post_init(&client, &url, Some("2025-11-25"), "2025-11-25").await; + assert_eq!(response.status(), 200); + + let response = post_init(&client, &url, Some("2025-03-26"), "2025-11-25").await; + assert_eq!(response.status(), 400); + let payload: serde_json::Value = response.json().await?; + assert_eq!(payload["error"]["code"], -32600); + + ct.cancel(); + Ok(()) +} + +// `notifications/initialized` returns HTTP 202 (Accepted) without +// requiring any protocol metadata, even under the seam. +#[tokio::test] +async fn seam_opt_in_notifications_return_202() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server(modern_required_config()).await; + + let body = r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#; + let response = post_seam(&client, &url, body, None, &[]).await; + assert_eq!( + response.status(), + reqwest::StatusCode::ACCEPTED, + "notifications must surface HTTP 202 regardless of metadata" + ); + + ct.cancel(); + Ok(()) +} + +// `server/discover` is still subject to the new HTTP-header precheck. +#[tokio::test] +async fn seam_opt_in_discover_rejects_missing_header() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server(modern_required_config()).await; + + let with_meta = r#"{"jsonrpc":"2.0","id":8,"method":"server/discover","params":{"_meta":{"io.modelcontextprotocol/protocolVersion":"2025-11-25","io.modelcontextprotocol/clientInfo":{"name":"test","version":"1.0"},"io.modelcontextprotocol/clientCapabilities":{}}}}"#; + let without_meta = r#"{"jsonrpc":"2.0","id":9,"method":"server/discover","params":{}}"#; + for body in [with_meta, without_meta] { + let response = post_seam( + &client, + &url, + body, + None, + &[("Mcp-Method", "server/discover")], + ) + .await; + assert_eq!(response.status(), 400); + let payload: serde_json::Value = response.json().await?; + assert_eq!(payload["error"]["code"], -32020); + } + + let response = post_seam( + &client, + &url, + without_meta, + Some("2025-11-25"), + &[("Mcp-Method", "server/discover")], + ) + .await; + assert_eq!(response.status(), 400); + let payload: serde_json::Value = response.json().await?; + assert_eq!(payload["error"]["code"], -32602); + + let response = post_seam( + &client, + &url, + with_meta, + Some("2025-11-25"), + &[("Mcp-Method", "server/discover")], + ) + .await; + assert_eq!(response.status(), 200); + + ct.cancel(); + Ok(()) +} diff --git a/crates/rmcp/tests/test_streamable_http_session_store.rs b/crates/rmcp/tests/test_streamable_http_session_store.rs index 91e7702..0c8f6a8 100644 --- a/crates/rmcp/tests/test_streamable_http_session_store.rs +++ b/crates/rmcp/tests/test_streamable_http_session_store.rs @@ -73,7 +73,7 @@ fn make_service( ) -> StreamableHttpService { StreamableHttpService::new(|| Ok(Calculator::new()), Default::default(), { let mut cfg = StreamableHttpServerConfig::default(); - cfg.stateful_mode = true; + cfg.legacy_session_mode = true; cfg.sse_keep_alive = None; cfg.cancellation_token = ct.child_token(); cfg.session_store = Some(session_store); @@ -147,7 +147,7 @@ async fn test_session_state_deleted_from_store_on_delete() -> anyhow::Result<()> let service = StreamableHttpService::new(|| Ok(Calculator::new()), session_manager.clone(), { let mut cfg = StreamableHttpServerConfig::default(); - cfg.stateful_mode = true; + cfg.legacy_session_mode = true; cfg.sse_keep_alive = None; cfg.cancellation_token = ct.child_token(); cfg.session_store = Some(store.clone()); @@ -219,7 +219,7 @@ fn spawn_server( ) -> (std::net::SocketAddr, tokio::task::JoinHandle<()>) { let svc = StreamableHttpService::new(|| Ok(Calculator::new()), session_manager, { let mut cfg = StreamableHttpServerConfig::default(); - cfg.stateful_mode = true; + cfg.legacy_session_mode = true; cfg.sse_keep_alive = None; cfg.cancellation_token = ct.child_token(); cfg.session_store = session_store; diff --git a/crates/rmcp/tests/test_streamable_http_stale_session.rs b/crates/rmcp/tests/test_streamable_http_stale_session.rs index d96d83c..c9ea214 100644 --- a/crates/rmcp/tests/test_streamable_http_stale_session.rs +++ b/crates/rmcp/tests/test_streamable_http_stale_session.rs @@ -5,15 +5,25 @@ not(feature = "local") ))] -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::{HashMap, VecDeque}, + sync::Arc, +}; +use futures::stream; +use http::{HeaderName, HeaderValue}; use rmcp::{ ServiceError, ServiceExt, - model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId}, + model::{ + CallToolRequestParams, ClientConfig, ClientJsonRpcMessage, ClientRequest, ErrorCode, + ErrorData, InitializeResult, PingRequest, ProtocolVersion, RequestId, ServerCapabilities, + ServerJsonRpcMessage, ServerResult, + }, transport::{ StreamableHttpClientTransport, streamable_http_client::{ StreamableHttpClient, StreamableHttpClientTransportConfig, StreamableHttpError, + StreamableHttpPostResponse, }, streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, @@ -25,6 +35,229 @@ use tokio_util::sync::CancellationToken; mod common; use common::calculator::Calculator; +#[derive(Debug, thiserror::Error)] +#[error("mock streamable http client error")] +struct MockClientError; + +#[derive(Clone)] +struct ReinitDropsAcceptedResponseClient { + state: Arc>, + stale_stream_cancelled: CancellationToken, + initial_request_accepted: Arc, + final_retry_accepted: Arc, +} + +struct MockState { + session_counter: usize, + posts: VecDeque, +} + +enum MockPost { + Initialize, + Initialized, + Accepted, + SessionExpired, +} + +impl ReinitDropsAcceptedResponseClient { + fn new() -> Self { + Self { + state: Arc::new(tokio::sync::Mutex::new(MockState { + session_counter: 0, + posts: VecDeque::from([ + MockPost::Initialize, + MockPost::Initialized, + MockPost::Accepted, + MockPost::SessionExpired, + MockPost::Initialize, + MockPost::Initialized, + MockPost::Accepted, + ]), + })), + stale_stream_cancelled: CancellationToken::new(), + initial_request_accepted: Arc::new(tokio::sync::Semaphore::new(0)), + final_retry_accepted: Arc::new(tokio::sync::Semaphore::new(0)), + } + } +} + +impl StreamableHttpClient for ReinitDropsAcceptedResponseClient { + type Error = MockClientError; + + async fn post_message( + &self, + _uri: Arc, + message: ClientJsonRpcMessage, + _session_id: Option>, + _auth_header: Option, + custom_headers: HashMap, + ) -> Result> { + let mut state = self.state.lock().await; + match state + .posts + .pop_front() + .expect("unexpected mock post_message call") + { + MockPost::Initialize => { + state.session_counter += 1; + let protocol_version = if state.session_counter == 1 { + ProtocolVersion::V_2025_11_25 + } else { + ProtocolVersion::V_2025_06_18 + }; + let id = match message { + ClientJsonRpcMessage::Request(request) => request.id, + other => panic!("expected initialize request, got {other:?}"), + }; + Ok(StreamableHttpPostResponse::Json( + ServerJsonRpcMessage::response( + ServerResult::InitializeResult( + InitializeResult::new( + ServerCapabilities::builder().enable_tools().build(), + ) + .with_protocol_version(protocol_version), + ), + id, + ), + Some(format!("session-{}", state.session_counter)), + )) + } + MockPost::Initialized => { + assert!( + matches!(message, ClientJsonRpcMessage::Notification(_)), + "expected initialized notification, got {message:?}" + ); + Ok(StreamableHttpPostResponse::Accepted) + } + MockPost::Accepted => { + if state.posts.is_empty() { + assert_eq!( + custom_headers + .get(&HeaderName::from_static("mcp-protocol-version")) + .and_then(|value| value.to_str().ok()), + Some(ProtocolVersion::V_2025_06_18.as_str()) + ); + self.final_retry_accepted.add_permits(1); + } else { + self.initial_request_accepted.add_permits(1); + } + Ok(StreamableHttpPostResponse::Accepted) + } + MockPost::SessionExpired => Err(StreamableHttpError::SessionExpired), + } + } + + async fn delete_session( + &self, + _uri: Arc, + _session_id: Arc, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + Ok(()) + } + + async fn get_stream( + &self, + _uri: Arc, + session_id: Option>, + _last_event_id: Option, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result< + futures::stream::BoxStream<'static, Result>, + StreamableHttpError, + > { + if session_id.as_deref() == Some("session-1") { + let cancel = self.stale_stream_cancelled.clone(); + Ok(Box::pin(stream::once(async move { + cancel.cancelled_owned().await; + Ok(sse_stream::Sse { + event: None, + data: Some( + serde_json::to_string(&ServerJsonRpcMessage::error( + ErrorData::new( + ErrorCode::INTERNAL_ERROR, + "stale stream should not deliver after re-init", + None, + ), + Some(RequestId::Number(2)), + )) + .expect("serialize stale error"), + ), + id: None, + retry: None, + }) + }))) + } else { + Ok(Box::pin(stream::pending())) + } + } +} + +#[tokio::test] +async fn test_reinitialization_completes_accepted_sse_request_instead_of_hanging() +-> anyhow::Result<()> { + let mock_client = ReinitDropsAcceptedResponseClient::new(); + let initial_request_accepted = mock_client.initial_request_accepted.clone(); + let final_retry_accepted = mock_client.final_retry_accepted.clone(); + let transport = StreamableHttpClientTransport::with_client( + mock_client, + StreamableHttpClientTransportConfig::with_uri("mock://mcp"), + ); + let mut client = ClientConfig::default().serve(transport).await?; + + let peer = client.peer().clone(); + let pending_call = tokio::spawn(async move { + peer.call_tool(CallToolRequestParams::new("slow_tool")) + .await + }); + + let _initial_permit = tokio::time::timeout( + std::time::Duration::from_secs(1), + initial_request_accepted.acquire(), + ) + .await + .expect("initial accepted request should be observed") + .expect("initial accepted request semaphore should stay open"); + + let reinit_trigger = { + let peer = client.peer().clone(); + tokio::spawn(async move { peer.list_tools(None).await }) + }; + + let _retry_permit = tokio::time::timeout( + std::time::Duration::from_secs(1), + final_retry_accepted.acquire(), + ) + .await + .expect("re-initialization retry should be accepted") + .expect("re-initialization retry semaphore should stay open"); + + let err = tokio::time::timeout(std::time::Duration::from_millis(100), pending_call) + .await + .expect("accepted SSE-backed request should complete instead of hanging")? + .expect_err( + "accepted request should fail after re-initialization drops its response stream", + ); + + match err { + ServiceError::McpError(error) => { + assert_eq!(error.code, ErrorCode::INTERNAL_ERROR); + assert!( + error.message.contains("session"), + "expected session-related error, got: {error}" + ); + } + other => panic!("expected McpError for orphaned request, got: {other:?}"), + } + + reinit_trigger.abort(); + let _ = client.close().await; + + Ok(()) +} + #[tokio::test] async fn test_stale_session_id_returns_status_aware_error() -> anyhow::Result<()> { let ct = CancellationToken::new(); diff --git a/crates/rmcp/tests/test_streamable_http_standard_headers.rs b/crates/rmcp/tests/test_streamable_http_standard_headers.rs new file mode 100644 index 0000000..9e70805 --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_standard_headers.rs @@ -0,0 +1,315 @@ +#![cfg(not(feature = "local"))] +//! SEP-2243 server-side validation of `Mcp-Method` / `Mcp-Name` / `Mcp-Param-*` headers. +use std::sync::Arc; + +use rmcp::{ + ServerHandler, + model::{ServerCapabilities, ServerConfig, Tool}, + transport::streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, + }, +}; +use tokio_util::sync::CancellationToken; + +const SEP_VERSION: &str = "2026-07-28"; + +/// Server exposing one tool whose `region` argument is promoted to `Mcp-Param-Region`. +#[derive(Clone, Default)] +struct HeaderValidationServer; + +impl ServerHandler for HeaderValidationServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + + fn get_tool(&self, name: &str) -> Option { + if name != "deploy" { + return None; + } + let schema = serde_json::json!({ + "type": "object", + "properties": { "region": { "type": "string", "x-mcp-header": "Region" } } + }); + let schema = schema.as_object().expect("object schema").clone(); + Some(Tool::new("deploy", "deploy a thing", Arc::new(schema))) + } +} + +async fn spawn_server() -> (reqwest::Client, String, CancellationToken) { + let config = StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_json_response(true) + .with_sse_keep_alive(None) + .with_cancellation_token(CancellationToken::new()); + let ct = config.cancellation_token.clone(); + let service: StreamableHttpService = + StreamableHttpService::new(|| Ok(HeaderValidationServer), Default::default(), config); + + let router = axum::Router::new().nest_service("/mcp", service); + let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = tcp_listener.local_addr().unwrap(); + tokio::spawn({ + let ct = ct.clone(); + async move { + let _ = axum::serve(tcp_listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + (reqwest::Client::new(), format!("http://{addr}/mcp"), ct) +} + +/// POSTs a `tools/call` with the given protocol-version and optional SEP-2243 headers. +async fn post_tool_call( + client: &reqwest::Client, + url: &str, + version: &str, + tool_name: &str, + arguments: serde_json::Value, + mcp_method: Option<&str>, + mcp_name: Option<&str>, + param_region: Option<&str>, +) -> reqwest::Response { + let body = serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": tool_name, + "arguments": arguments, + } + }); + let mut req = client + .post(url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", version) + .body(body.to_string()); + if let Some(method) = mcp_method { + req = req.header("Mcp-Method", method); + } + if let Some(name) = mcp_name { + req = req.header("Mcp-Name", name); + } + if let Some(region) = param_region { + req = req.header("Mcp-Param-Region", region); + } + req.send().await.expect("send tools/call request") +} + +#[tokio::test] +async fn accepts_matching_standard_headers() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server().await; + + // Matching headers pass validation and reach dispatch. (Stateless mode without a + // prior initialize yields an unrelated -32601, which still proves -32020 was not raised.) + let response = post_tool_call( + &client, + &url, + SEP_VERSION, + "sum", + serde_json::json!({ "a": 1, "b": 2 }), + Some("tools/call"), + Some("sum"), + None, + ) + .await; + let body: serde_json::Value = response.json().await?; + assert_ne!( + body["error"]["code"], -32020, + "matching headers must not be rejected as a header mismatch, got: {body}" + ); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn rejects_method_mismatch_with_32020() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server().await; + + let response = post_tool_call( + &client, + &url, + SEP_VERSION, + "sum", + serde_json::json!({ "a": 1, "b": 2 }), + Some("tools/list"), + Some("sum"), + None, + ) + .await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32020); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn rejects_missing_method_header_with_32020() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server().await; + + let response = post_tool_call( + &client, + &url, + SEP_VERSION, + "sum", + serde_json::json!({ "a": 1, "b": 2 }), + None, + Some("sum"), + None, + ) + .await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32020); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn rejects_name_mismatch_with_32020() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server().await; + + let response = post_tool_call( + &client, + &url, + SEP_VERSION, + "sum", + serde_json::json!({ "a": 1, "b": 2 }), + Some("tools/call"), + Some("product"), + None, + ) + .await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32020); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn skips_validation_for_pre_sep_version() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server().await; + + // Older version: headers are not enforced even when absent. + let response = post_tool_call( + &client, + &url, + "2025-11-25", + "sum", + serde_json::json!({ "a": 1, "b": 2 }), + None, + None, + None, + ) + .await; + let body: serde_json::Value = response.json().await?; + assert_ne!( + body["error"]["code"], -32020, + "pre-SEP versions must skip header validation, got: {body}" + ); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn accepts_matching_param_header() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server().await; + + let response = post_tool_call( + &client, + &url, + SEP_VERSION, + "deploy", + serde_json::json!({ "region": "us-west1" }), + Some("tools/call"), + Some("deploy"), + Some("us-west1"), + ) + .await; + let body: serde_json::Value = response.json().await?; + assert_ne!( + body["error"]["code"], -32020, + "matching Mcp-Param-* must not be rejected, got: {body}" + ); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn rejects_param_mismatch_with_32020() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server().await; + + let response = post_tool_call( + &client, + &url, + SEP_VERSION, + "deploy", + serde_json::json!({ "region": "us-west1" }), + Some("tools/call"), + Some("deploy"), + Some("eu-central1"), + ) + .await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32020); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn rejects_decoded_base64_param_mismatch_with_32020() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server().await; + + let response = post_tool_call( + &client, + &url, + SEP_VERSION, + "deploy", + serde_json::json!({ "region": "us-west1" }), + Some("tools/call"), + Some("deploy"), + Some("=?base64?ZXUtY2VudHJhbDE=?="), + ) + .await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32020); + + ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn rejects_missing_param_header_with_32020() -> anyhow::Result<()> { + let (client, url, ct) = spawn_server().await; + + // `region` argument is present but the annotated `Mcp-Param-Region` header is absent. + let response = post_tool_call( + &client, + &url, + SEP_VERSION, + "deploy", + serde_json::json!({ "region": "us-west1" }), + Some("tools/call"), + Some("deploy"), + None, + ) + .await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32020); + + ct.cancel(); + Ok(()) +} diff --git a/crates/rmcp/tests/test_structured_output.rs b/crates/rmcp/tests/test_structured_output.rs index bb0d5e0..f3416e5 100644 --- a/crates/rmcp/tests/test_structured_output.rs +++ b/crates/rmcp/tests/test_structured_output.rs @@ -3,7 +3,7 @@ use rmcp::{ Json, ServerHandler, handler::server::{router::tool::ToolRouter, tool::IntoCallToolResult, wrapper::Parameters}, - model::{CallToolResult, ContentBlock, ServerResult, Tool}, + model::{CallToolResponse, CallToolResult, ContentBlock, ServerResult, Tool}, tool, tool_handler, tool_router, }; use schemars::JsonSchema; @@ -93,6 +93,27 @@ impl TestServer { Err("User not found".to_string()) } } + + /// Tool that returns a list of calculation results + #[tool( + name = "calculate-list", + description = "Return a list of calculation results" + )] + pub async fn calculate_list( + &self, + params: Parameters, + ) -> Result>, String> { + Ok(Json(vec![CalculationResult { + sum: params.0.a + params.0.b, + product: params.0.a * params.0.b, + }])) + } + + /// Tool that returns a count + #[tool(name = "get-count", description = "Return a count")] + pub async fn get_count(&self) -> Result, String> { + Ok(Json(42)) + } } #[tokio::test] @@ -224,11 +245,13 @@ async fn test_structured_return_conversion() { }; let structured = Json(calc_result); - let result: Result = + let result: Result = rmcp::handler::server::tool::IntoCallToolResult::into_call_tool_result(structured); assert!(result.is_ok()); - let call_result = result.unwrap(); + let CallToolResponse::Complete(call_result) = result.unwrap() else { + panic!("expected complete CallToolResult"); + }; // Tools which return structured content should also return a serialized version as // Content::text for backwards compatibility. @@ -285,11 +308,13 @@ async fn test_output_schema_requires_structured_content() { let result = server.calculate(params).await.unwrap(); // Convert the Json to CallToolResult - let call_result: Result = + let call_result: Result = IntoCallToolResult::into_call_tool_result(result); assert!(call_result.is_ok()); - let call_result = call_result.unwrap(); + let CallToolResponse::Complete(call_result) = call_result.unwrap() else { + panic!("expected complete CallToolResult"); + }; // Verify it has structured_content and content assert!(call_result.structured_content.is_some()); @@ -360,3 +385,38 @@ fn test_call_tool_result_deserialize_without_content() { assert!(result.content.is_empty()); assert!(result.structured_content.is_some()); } + +#[tokio::test] +async fn test_tool_with_array_output_schema() { + let server = TestServer::new(); + let tools = server.tool_router.list_all(); + + // Find the calculate-list tool + let calculate_list_tool = tools.iter().find(|t| t.name == "calculate-list").unwrap(); + + // Verify it has an output schema + assert!(calculate_list_tool.output_schema.is_some()); + + let schema = calculate_list_tool.output_schema.as_ref().unwrap(); + + // Check that the schema contains array type + let schema_str = serde_json::to_string(schema).unwrap(); + assert!(schema_str.contains("array")); +} + +#[tokio::test] +async fn test_tool_with_primitive_output_schema() { + let server = TestServer::new(); + let tools = server.tool_router.list_all(); + + // Find the get-count tool + let get_count_tool = tools.iter().find(|t| t.name == "get-count").unwrap(); + + // Verify it has an output schema + assert!(get_count_tool.output_schema.is_some()); + + let schema = get_count_tool.output_schema.as_ref().unwrap(); + + // Check that the schema contains integer type + assert_eq!(schema.get("type"), Some(&serde_json::json!("integer"))); +} diff --git a/crates/rmcp/tests/test_subscriptions.rs b/crates/rmcp/tests/test_subscriptions.rs new file mode 100644 index 0000000..7359530 --- /dev/null +++ b/crates/rmcp/tests/test_subscriptions.rs @@ -0,0 +1,727 @@ +#![cfg(all( + not(feature = "local"), + feature = "client", + feature = "server", + feature = "transport-io" +))] + +use std::{ + num::NonZeroUsize, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use rmcp::{ + ClientHandler, ClientServiceExt, ServerHandler, ServiceExt, + model::{ + ClientNotification, ClientRequest, DiscoverResult, GetMeta, Implementation, + NotificationMetaObject, PromptListChangedNotification, ProtocolVersion, ServerCapabilities, + ServerConfig, ServerNotification, ServerResult, SubscriptionFilter, + SubscriptionsAcknowledgedNotification, SubscriptionsAcknowledgedNotificationParams, + SubscriptionsListenResult, + }, + service::{ + NotificationContext, RequestContext, RoleClient, RoleServer, SubscriptionContext, + SubscriptionEnd, SubscriptionSendError, SubscriptionSink, + }, +}; +use tokio::sync::{Mutex, Notify}; + +struct ToolsOnlyServer; + +#[derive(Clone)] +struct CountingClient { + tool_changes: Arc, +} + +impl ClientHandler for CountingClient { + async fn on_tool_list_changed(&self, _context: NotificationContext) { + self.tool_changes.fetch_add(1, Ordering::Relaxed); + } +} + +impl ServerHandler for ToolsOnlyServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tool_list_changed() + .build(), + ) + .with_server_info(Implementation::new("tools-only-server", "1.0.0")) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.supported_by(&self.get_info().capabilities)) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + context + .sink() + .notify_tool_list_changed() + .await + .expect("accepted tool notification"); + assert!(matches!( + context.sink().notify_prompt_list_changed().await, + Err(SubscriptionSendError::NotificationNotAccepted( + "notifications/prompts/list_changed" + )) + )); + Ok(()) + } +} + +struct ToolsAndPromptsServer; + +impl ServerHandler for ToolsAndPromptsServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tool_list_changed() + .enable_prompts() + .enable_prompts_list_changed() + .build(), + ) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.supported_by(&self.get_info().capabilities)) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + if context.accepted().tools_list_changed == Some(true) { + context + .sink() + .notify_tool_list_changed() + .await + .expect("send tool notification"); + } + if context.accepted().prompts_list_changed == Some(true) { + context + .sink() + .notify_prompt_list_changed() + .await + .expect("send prompt notification"); + } + Ok(()) + } +} + +struct ResourceSubscriptionServer; + +impl ServerHandler for ResourceSubscriptionServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_resources() + .enable_resources_subscribe() + .build(), + ) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.supported_by(&self.get_info().capabilities)) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + context + .sink() + .notify_resource_updated("file:///accepted") + .await + .expect("accepted URI"); + assert!(matches!( + context + .sink() + .notify_resource_updated("file:///not-requested") + .await, + Err(SubscriptionSendError::NotificationNotAccepted( + "notifications/resources/updated" + )) + )); + Ok(()) + } +} + +struct RemoteCancellationServer; + +impl ServerHandler for RemoteCancellationServer { + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.clone()) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + context + .request_context() + .peer + .notify_cancelled(rmcp::model::CancelledNotificationParam::new( + Some(context.sink().id().clone()), + Some("server shutdown".to_owned()), + )) + .await + .expect("send server cancellation"); + std::future::pending().await + } +} + +struct FloodServer; + +impl ServerHandler for FloodServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tool_list_changed() + .build(), + ) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.clone()) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + for _ in 0..10 { + if context.sink().notify_tool_list_changed().await.is_err() { + break; + } + } + context.cancelled().await; + Ok(()) + } +} + +#[derive(Clone)] +struct ClosedSinkServer { + sink: Arc>>, +} + +struct LeakyServer; + +impl ServerHandler for LeakyServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tool_list_changed() + .build(), + ) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.supported_by(&self.get_info().capabilities)) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + let mut notification = + ServerNotification::PromptListChangedNotification(PromptListChangedNotification { + method: Default::default(), + extensions: Default::default(), + }); + notification + .get_meta_mut() + .set_subscription_id(context.sink().id().clone()); + context + .request_context() + .peer + .send_notification(notification) + .await + .expect("send deliberately invalid notification"); + std::future::pending().await + } +} + +struct MalformedAcknowledgmentServer { + cancelled: Arc, +} + +impl rmcp::service::Service for MalformedAcknowledgmentServer { + async fn handle_request( + &self, + request: ClientRequest, + context: RequestContext, + ) -> Result { + match request { + ClientRequest::DiscoverRequest(_) => Ok(ServerResult::DiscoverResult( + DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + ServerCapabilities::builder() + .enable_tools() + .enable_tool_list_changed() + .enable_prompts() + .enable_prompts_list_changed() + .build(), + ) + .with_server_info(Implementation::new("malformed-ack-server", "1.0.0")), + )), + ClientRequest::SubscriptionsListenRequest(_) => { + let mut acknowledgment = SubscriptionsAcknowledgedNotification::new( + SubscriptionsAcknowledgedNotificationParams::new( + SubscriptionFilter::builder().prompts_list_changed().build(), + ), + ); + let mut meta = NotificationMetaObject::new(); + meta.set_subscription_id(context.id.clone()); + acknowledgment.extensions.insert(meta); + context + .peer + .send_notification(ServerNotification::SubscriptionsAcknowledgedNotification( + acknowledgment, + )) + .await + .map_err(|error| rmcp::ErrorData::internal_error(error.to_string(), None))?; + context.ct.cancelled().await; + Ok(ServerResult::SubscriptionsListenResult( + SubscriptionsListenResult::complete(context.id), + )) + } + _ => Err(rmcp::ErrorData::invalid_request( + "unexpected test request", + None, + )), + } + } + + async fn handle_notification( + &self, + notification: ClientNotification, + _context: NotificationContext, + ) -> Result<(), rmcp::ErrorData> { + if matches!(notification, ClientNotification::CancelledNotification(_)) { + self.cancelled.notify_one(); + } + Ok(()) + } + + fn get_info(&self) -> rmcp::model::ServerConfig { + ServerConfig::default() + } +} + +impl ServerHandler for ClosedSinkServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tool_list_changed() + .build(), + ) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.supported_by(&self.get_info().capabilities)) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + self.sink.lock().await.replace(context.sink().clone()); + Ok(()) + } +} + +#[derive(Clone)] +struct CancellationServer { + cancelled: Arc, +} + +impl ServerHandler for CancellationServer { + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.clone()) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + context.cancelled().await; + self.cancelled.notify_one(); + Ok(()) + } +} + +#[derive(Clone)] +struct AbruptServer { + started: Arc, +} + +impl ServerHandler for AbruptServer { + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.clone()) + } + + async fn listen(&self, _context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + self.started.notify_one(); + std::future::pending().await + } +} + +async fn modern_client( + server: S, +) -> anyhow::Result> { + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + tokio::spawn(async move { + let server = server.serve(server_transport).await?; + server.waiting().await?; + anyhow::Ok(()) + }); + ().serve_with_lifecycle( + client_transport, + rmcp::ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .map_err(Into::into) +} + +#[tokio::test] +async fn listen_exposes_acknowledged_filter_and_graceful_result() -> anyhow::Result<()> { + let client = modern_client(ToolsOnlyServer).await?; + let mut subscription = client + .listen( + SubscriptionFilter::builder() + .tools_list_changed() + .prompts_list_changed() + .build(), + ) + .await?; + + assert_eq!( + subscription.acknowledged(), + &SubscriptionFilter::builder().tools_list_changed().build() + ); + + let notification = tokio::time::timeout(Duration::from_secs(5), subscription.next()) + .await?? + .expect("tool notification"); + assert!(matches!( + notification, + ServerNotification::ToolListChangedNotification(_) + )); + assert_eq!( + notification.get_meta().subscription_id(), + Some(subscription.id().clone()) + ); + + assert!( + tokio::time::timeout(Duration::from_secs(5), subscription.next()) + .await?? + .is_none() + ); + let Some(SubscriptionEnd::Graceful(result)) = subscription.end() else { + panic!("expected graceful final result"); + }; + assert_eq!( + result.meta.subscription_id().as_ref(), + Some(subscription.id()) + ); + assert_eq!( + result + .meta + .server_info() + .expect("graceful result should contain valid server info"), + Implementation::new("tools-only-server", "1.0.0") + ); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn typed_subscription_notifications_do_not_reach_handler_callbacks() -> anyhow::Result<()> { + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + tokio::spawn(async move { + let server = ToolsOnlyServer.serve(server_transport).await?; + server.waiting().await?; + anyhow::Ok(()) + }); + let tool_changes = Arc::new(AtomicUsize::new(0)); + let client = CountingClient { + tool_changes: tool_changes.clone(), + } + .serve_with_lifecycle( + client_transport, + rmcp::ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await?; + let mut subscription = client + .listen(SubscriptionFilter::builder().tools_list_changed().build()) + .await?; + + assert!(subscription.next().await?.is_some()); + assert!(subscription.next().await?.is_none()); + tokio::task::yield_now().await; + assert_eq!(tool_changes.load(Ordering::Relaxed), 0); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn discover_lifecycle_allows_subscriptions_with_older_application_version() +-> anyhow::Result<()> { + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + tokio::spawn(async move { + let server = ToolsOnlyServer.serve(server_transport).await?; + server.waiting().await?; + anyhow::Ok(()) + }); + let client = () + .serve_with_lifecycle( + client_transport, + rmcp::ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2025_11_25], + }, + ) + .await?; + let mut subscription = client + .listen(SubscriptionFilter::builder().tools_list_changed().build()) + .await?; + + assert!(subscription.next().await?.is_some()); + assert!(subscription.next().await?.is_none()); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn concurrent_subscriptions_are_demultiplexed_by_request_id() -> anyhow::Result<()> { + let client = modern_client(ToolsAndPromptsServer).await?; + let (tools, prompts) = tokio::join!( + client.listen(SubscriptionFilter::builder().tools_list_changed().build()), + client.listen(SubscriptionFilter::builder().prompts_list_changed().build()) + ); + let mut tools = tools?; + let mut prompts = prompts?; + + let tool_notification = tools.next().await?.expect("tool notification"); + let prompt_notification = prompts.next().await?.expect("prompt notification"); + assert!(matches!( + tool_notification, + ServerNotification::ToolListChangedNotification(_) + )); + assert!(matches!( + prompt_notification, + ServerNotification::PromptListChangedNotification(_) + )); + assert_ne!(tools.id(), prompts.id()); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn stdio_cancellation_sends_cancelled_for_the_listen_request() -> anyhow::Result<()> { + let cancelled = Arc::new(Notify::new()); + let client = modern_client(CancellationServer { + cancelled: cancelled.clone(), + }) + .await?; + let mut subscription = client.listen(SubscriptionFilter::new()).await?; + + subscription.cancel().await?; + assert!(matches!( + subscription.end(), + Some(SubscriptionEnd::Cancelled) + )); + tokio::time::timeout(Duration::from_secs(5), cancelled.notified()).await?; + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn resource_updates_are_filtered_by_exact_uri_membership() -> anyhow::Result<()> { + let client = modern_client(ResourceSubscriptionServer).await?; + let mut subscription = client + .listen( + SubscriptionFilter::builder() + .resource_subscription("file:///accepted") + .build(), + ) + .await?; + + let notification = subscription.next().await?.expect("resource update"); + let ServerNotification::ResourceUpdatedNotification(update) = notification else { + panic!("expected resource update"); + }; + assert_eq!(update.params.uri, "file:///accepted"); + assert!(subscription.next().await?.is_none()); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn stdio_server_cancellation_ends_the_matching_subscription() -> anyhow::Result<()> { + let client = modern_client(RemoteCancellationServer).await?; + let mut subscription = client.listen(SubscriptionFilter::new()).await?; + + assert!(subscription.next().await?.is_none()); + assert!(matches!( + subscription.end(), + Some(SubscriptionEnd::Cancelled) + )); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn slow_consumer_reports_subscription_lag() -> anyhow::Result<()> { + let client = modern_client(FloodServer).await?; + let mut subscription = client + .listen_with_capacity( + SubscriptionFilter::builder().tools_list_changed().build(), + NonZeroUsize::MIN, + ) + .await?; + + tokio::time::sleep(Duration::from_millis(50)).await; + while subscription.next().await?.is_some() {} + assert!( + matches!( + subscription.end(), + Some(SubscriptionEnd::Lagged { capacity: 1 }) + ), + "unexpected end: {:?}", + subscription.end() + ); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn sink_rejects_notifications_after_graceful_completion() -> anyhow::Result<()> { + let sink = Arc::new(Mutex::new(None)); + let client = modern_client(ClosedSinkServer { sink: sink.clone() }).await?; + let mut subscription = client + .listen(SubscriptionFilter::builder().tools_list_changed().build()) + .await?; + assert!(subscription.next().await?.is_none()); + + let sink = sink.lock().await.clone().expect("captured sink"); + assert!(matches!( + sink.notify_tool_list_changed().await, + Err(SubscriptionSendError::SubscriptionClosed) + )); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn client_rejects_notifications_outside_the_acknowledged_filter() -> anyhow::Result<()> { + let client = modern_client(LeakyServer).await?; + let mut subscription = client + .listen(SubscriptionFilter::builder().tools_list_changed().build()) + .await?; + + assert!(matches!( + subscription.next().await, + Err(rmcp::ServiceError::UnexpectedResponse) + )); + assert!(matches!(subscription.end(), Some(SubscriptionEnd::Abrupt))); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn malformed_acknowledgment_cancels_pending_listen_request() -> anyhow::Result<()> { + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + let cancelled = Arc::new(Notify::new()); + let server_cancelled = cancelled.clone(); + tokio::spawn(async move { + let server = MalformedAcknowledgmentServer { + cancelled: server_cancelled, + } + .serve(server_transport) + .await?; + server.waiting().await?; + anyhow::Ok(()) + }); + let client = () + .serve_with_lifecycle( + client_transport, + rmcp::ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await?; + + assert!(matches!( + client + .listen(SubscriptionFilter::builder().tools_list_changed().build()) + .await, + Err(rmcp::ServiceError::UnexpectedResponse) + )); + tokio::time::timeout(Duration::from_secs(5), cancelled.notified()).await?; + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn transport_close_without_final_result_is_abrupt() -> anyhow::Result<()> { + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + let started = Arc::new(Notify::new()); + let server_started = started.clone(); + tokio::spawn(async move { + let server = AbruptServer { + started: server_started.clone(), + } + .serve(server_transport) + .await?; + server_started.notified().await; + server.cancel().await?; + anyhow::Ok(()) + }); + let client = () + .serve_with_lifecycle( + client_transport, + rmcp::ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await?; + let mut subscription = client.listen(SubscriptionFilter::new()).await?; + + assert!( + tokio::time::timeout(Duration::from_secs(5), subscription.next()) + .await?? + .is_none() + ); + assert!(matches!(subscription.end(), Some(SubscriptionEnd::Abrupt))); + Ok(()) +} diff --git a/crates/rmcp/tests/test_subscriptions_model.rs b/crates/rmcp/tests/test_subscriptions_model.rs new file mode 100644 index 0000000..8d79fb1 --- /dev/null +++ b/crates/rmcp/tests/test_subscriptions_model.rs @@ -0,0 +1,254 @@ +use rmcp::model::{ + ClientJsonRpcMessage, ClientRequest, GetMeta, Implementation, JsonRpcNotification, + JsonRpcRequest, NotificationMetaObject, RequestId, RequestMetaObject, ServerJsonRpcMessage, + ServerNotification, SubscriptionFilter, SubscriptionsAcknowledgedNotification, + SubscriptionsAcknowledgedNotificationParams, SubscriptionsListenRequest, + SubscriptionsListenRequestParams, SubscriptionsListenResult, SubscriptionsListenResultMeta, +}; +use serde_json::json; + +#[test] +fn subscription_filter_serializes_only_opted_in_notifications() { + let filter = SubscriptionFilter::builder() + .tools_list_changed() + .resource_subscription("file:///one") + .resource_subscription("file:///two") + .build(); + + assert_eq!( + serde_json::to_value(filter).expect("serialize filter"), + json!({ + "toolsListChanged": true, + "resourceSubscriptions": ["file:///one", "file:///two"], + }) + ); +} + +#[test] +fn subscription_filter_subset_is_order_independent_and_ignores_false_flags() { + let requested = SubscriptionFilter::builder() + .tools_list_changed() + .resource_subscriptions(["file:///one", "file:///two"]) + .build(); + let mut accepted = SubscriptionFilter::builder() + .resource_subscriptions(["file:///two", "file:///one"]) + .build(); + accepted.tools_list_changed = Some(false); + + assert!(accepted.is_subset_of(&requested)); +} + +#[test] +fn subscription_filter_omits_empty_resource_intersection() { + let requested = SubscriptionFilter::builder() + .resource_subscription("file:///requested") + .build(); + let accepted = SubscriptionFilter::builder() + .resource_subscription("file:///different") + .build(); + + assert_eq!( + serde_json::to_value(requested.intersection(&accepted)).expect("serialize intersection"), + json!({}) + ); +} + +#[test] +fn listen_request_round_trips_required_fields_and_arbitrary_metadata() { + let mut request = SubscriptionsListenRequest::new(SubscriptionsListenRequestParams::new( + SubscriptionFilter::builder().prompts_list_changed().build(), + )); + let mut meta = RequestMetaObject::new(); + meta.insert("com.example/request".into(), json!("value")); + request.extensions.insert(meta); + let message = ClientJsonRpcMessage::request( + ClientRequest::SubscriptionsListenRequest(request), + RequestId::String("subscription-1".into()), + ); + + let value = serde_json::to_value(&message).expect("serialize listen request"); + assert_eq!( + value, + json!({ + "jsonrpc": "2.0", + "id": "subscription-1", + "method": "subscriptions/listen", + "params": { + "_meta": { + "com.example/request": "value", + }, + "notifications": { + "promptsListChanged": true, + }, + }, + }) + ); + + let round_trip: ClientJsonRpcMessage = + serde_json::from_value(value).expect("deserialize listen request"); + let ClientJsonRpcMessage::Request(JsonRpcRequest { request, .. }) = round_trip else { + panic!("expected request"); + }; + let ClientRequest::SubscriptionsListenRequest(request) = request else { + panic!("expected subscriptions/listen request"); + }; + assert_eq!( + request + .extensions + .get::() + .and_then(|meta| meta.get("com.example/request")), + Some(&json!("value")) + ); +} + +#[test] +fn acknowledged_notification_round_trips_numeric_subscription_id_and_metadata() { + let mut notification = SubscriptionsAcknowledgedNotification::new( + SubscriptionsAcknowledgedNotificationParams::new( + SubscriptionFilter::builder() + .resources_list_changed() + .build(), + ), + ); + let mut meta = NotificationMetaObject::new(); + meta.set_subscription_id(RequestId::Number(7)); + meta.insert("com.example/notification".into(), json!(true)); + notification.extensions.insert(meta); + let message = ServerJsonRpcMessage::notification( + ServerNotification::SubscriptionsAcknowledgedNotification(notification), + ); + + let value = serde_json::to_value(&message).expect("serialize acknowledgment"); + assert_eq!( + value, + json!({ + "jsonrpc": "2.0", + "method": "notifications/subscriptions/acknowledged", + "params": { + "_meta": { + "io.modelcontextprotocol/subscriptionId": 7, + "com.example/notification": true, + }, + "notifications": { + "resourcesListChanged": true, + }, + }, + }) + ); + + let round_trip: ServerJsonRpcMessage = + serde_json::from_value(value).expect("deserialize acknowledgment"); + let ServerJsonRpcMessage::Notification(JsonRpcNotification { notification, .. }) = round_trip + else { + panic!("expected notification"); + }; + assert_eq!( + notification.get_meta().subscription_id(), + Some(RequestId::Number(7)) + ); + assert_eq!( + notification.get_meta().get("com.example/notification"), + Some(&json!(true)) + ); +} + +#[test] +fn listen_result_requires_matching_string_subscription_id_and_preserves_metadata() { + let mut meta = SubscriptionsListenResultMeta::new(RequestId::String("subscription-2".into())); + meta.set_server_info(Implementation::new("test-server", "1.0.0")); + meta.insert("com.example/result".into(), json!({ "reason": "shutdown" })); + let result = SubscriptionsListenResult::new(meta); + + let value = serde_json::to_value(&result).expect("serialize listen result"); + assert_eq!( + value, + json!({ + "resultType": "complete", + "_meta": { + "io.modelcontextprotocol/subscriptionId": "subscription-2", + "io.modelcontextprotocol/serverInfo": { + "name": "test-server", + "version": "1.0.0", + }, + "com.example/result": { + "reason": "shutdown", + }, + }, + }) + ); + + let round_trip: SubscriptionsListenResult = + serde_json::from_value(value).expect("deserialize listen result"); + assert_eq!( + round_trip.meta.subscription_id(), + Some(RequestId::String("subscription-2".into())) + ); + assert_eq!( + round_trip.meta.server_info(), + Some(Implementation::new("test-server", "1.0.0")) + ); + assert_eq!( + round_trip.meta.get("com.example/result"), + Some(&json!({ "reason": "shutdown" })) + ); +} + +#[test] +fn listen_result_meta_returns_none_after_required_id_is_removed() { + let mut meta = SubscriptionsListenResultMeta::new(RequestId::Number(1)); + meta.remove("io.modelcontextprotocol/subscriptionId"); + + assert_eq!(meta.subscription_id(), None); +} + +#[cfg(feature = "schemars")] +#[test] +fn subscription_schemas_mark_only_draft_required_fields_as_required() { + let request_schema = + serde_json::to_value(schemars::schema_for!(SubscriptionsListenRequestParams)) + .expect("request schema"); + let filter_schema = + serde_json::to_value(schemars::schema_for!(SubscriptionFilter)).expect("filter schema"); + let acknowledgment_schema = serde_json::to_value(schemars::schema_for!( + SubscriptionsAcknowledgedNotificationParams + )) + .expect("acknowledgment schema"); + let result_schema = serde_json::to_value(schemars::schema_for!(SubscriptionsListenResult)) + .expect("result schema"); + + assert_eq!( + request_schema["required"], + json!(["_meta", "notifications"]) + ); + assert_eq!( + request_schema["properties"]["_meta"]["required"], + json!([ + "io.modelcontextprotocol/protocolVersion", + "io.modelcontextprotocol/clientCapabilities" + ]) + ); + assert!(filter_schema.get("required").is_none()); + assert_eq!( + filter_schema["properties"]["toolsListChanged"]["type"], + "boolean" + ); + assert_eq!( + filter_schema["properties"]["resourceSubscriptions"]["type"], + "array" + ); + assert_eq!(acknowledgment_schema["required"], json!(["notifications"])); + assert_eq!( + acknowledgment_schema["properties"]["_meta"]["$ref"], + "#/$defs/NotificationMetaObject" + ); + let result_meta_schema = &result_schema["$defs"]["SubscriptionsListenResultMeta"]; + assert_eq!( + result_meta_schema["properties"]["io.modelcontextprotocol/serverInfo"]["allOf"][0]["$ref"], + "#/$defs/Implementation" + ); + assert_eq!( + result_meta_schema["required"], + json!(["io.modelcontextprotocol/subscriptionId"]) + ); + assert_eq!(result_schema["required"], json!(["resultType", "_meta"])); +} diff --git a/crates/rmcp/tests/test_subscriptions_streamable_http.rs b/crates/rmcp/tests/test_subscriptions_streamable_http.rs new file mode 100644 index 0000000..70f102f --- /dev/null +++ b/crates/rmcp/tests/test_subscriptions_streamable_http.rs @@ -0,0 +1,318 @@ +#![cfg(all( + not(feature = "local"), + feature = "client", + feature = "server", + feature = "transport-streamable-http-client-reqwest", + feature = "transport-streamable-http-server" +))] + +use std::{ + borrow::Cow, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use rmcp::{ + ClientLifecycleMode, ClientServiceExt, ServerHandler, + model::{ + ClientConfig, ClientRequest, Implementation, ListToolsRequest, ProtocolVersion, + RequestMetaObject, ServerCapabilities, ServerConfig, ServerNotification, + SubscriptionFilter, + }, + service::{PeerRequestOptions, SubscriptionContext, SubscriptionEnd}, + transport::{ + StreamableHttpClientTransport, + streamable_http_client::StreamableHttpClientTransportConfig, + streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, + }, + }, +}; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +#[derive(Clone)] +struct HttpSubscriptionServer { + cancelled: Arc, + started: Arc, + ending: ServerEnding, +} + +#[derive(Clone, Copy)] +enum ServerEnding { + ClientCancellation, + Graceful, + Abrupt, +} + +impl ServerHandler for HttpSubscriptionServer { + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(&[ProtocolVersion::V_2026_07_28, ProtocolVersion::V_2025_11_25]) + } + + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tool_list_changed() + .build(), + ) + .with_server_info(Implementation::new("http-subscription-server", "1.0.0")) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.supported_by(&self.get_info().capabilities)) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), rmcp::ErrorData> { + context + .sink() + .notify_tool_list_changed() + .await + .expect("send tool notification"); + self.started.notify_one(); + match self.ending { + ServerEnding::Graceful => Ok(()), + ServerEnding::ClientCancellation => { + context.cancelled().await; + self.cancelled.notify_one(); + Ok(()) + } + ServerEnding::Abrupt => std::future::pending().await, + } + } +} + +async fn spawn_server( + ending: ServerEnding, +) -> ( + String, + CancellationToken, + Arc, + Arc, + Arc, +) { + let cancellation_token = CancellationToken::new(); + let subscription_cancelled = Arc::new(Notify::new()); + let subscription_started = Arc::new(Notify::new()); + let server = HttpSubscriptionServer { + cancelled: subscription_cancelled.clone(), + started: subscription_started.clone(), + ending, + }; + let service: StreamableHttpService = + StreamableHttpService::new( + move || Ok(server.clone()), + Default::default(), + StreamableHttpServerConfig::default() + .with_legacy_session_mode(true) + .with_json_response(true) + .with_sse_keep_alive(Some(Duration::from_millis(50))) + .with_cancellation_token(cancellation_token.child_token()), + ); + let get_requests = Arc::new(AtomicUsize::new(0)); + let observed_get_requests = get_requests.clone(); + let router = + axum::Router::new() + .nest_service("/mcp", service) + .layer(axum::middleware::from_fn( + move |request: axum::extract::Request, next: axum::middleware::Next| { + let observed_get_requests = observed_get_requests.clone(); + async move { + if request.method() == axum::http::Method::GET { + observed_get_requests.fetch_add(1, Ordering::Relaxed); + } + next.run(request).await + } + }, + )); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("listener"); + let address = listener.local_addr().expect("listener address"); + tokio::spawn({ + let cancellation_token = cancellation_token.clone(); + async move { + let _ = axum::serve(listener, router) + .with_graceful_shutdown(async move { cancellation_token.cancelled_owned().await }) + .await; + } + }); + ( + format!("http://{address}/mcp"), + cancellation_token, + subscription_cancelled, + subscription_started, + get_requests, + ) +} + +#[tokio::test] +async fn modern_http_listen_uses_post_stream_and_cancels_by_closing_it() -> anyhow::Result<()> { + let (url, server_ct, subscription_cancelled, _, get_requests) = + spawn_server(ServerEnding::ClientCancellation).await; + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(url.clone()), + ); + let client = ClientConfig::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await?; + assert_eq!(get_requests.load(Ordering::Relaxed), 0); + let mut subscription = client + .listen(SubscriptionFilter::builder().tools_list_changed().build()) + .await?; + + assert!(matches!( + subscription.next().await?.expect("tool notification"), + ServerNotification::ToolListChangedNotification(_) + )); + let mut older_version = RequestMetaObject::new(); + older_version.set_protocol_version(ProtocolVersion::V_2025_11_25); + client + .send_request_with_option( + ClientRequest::ListToolsRequest(ListToolsRequest { + method: Default::default(), + params: None, + extensions: Default::default(), + }), + PeerRequestOptions::no_options().with_meta(older_version), + ) + .await? + .await_response() + .await?; + subscription.cancel().await?; + tokio::time::timeout(Duration::from_secs(5), subscription_cancelled.notified()).await?; + + client.cancel().await?; + server_ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn modern_http_graceful_close_returns_final_listen_result() -> anyhow::Result<()> { + let (url, server_ct, _, _, _) = spawn_server(ServerEnding::Graceful).await; + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(url), + ); + let client = ClientConfig::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await?; + let mut subscription = client + .listen(SubscriptionFilter::builder().tools_list_changed().build()) + .await?; + + assert!(subscription.next().await?.is_some()); + assert!(subscription.next().await?.is_none()); + let Some(SubscriptionEnd::Graceful(result)) = subscription.end() else { + panic!("expected graceful final result"); + }; + assert_eq!( + result + .meta + .server_info() + .expect("graceful result should contain valid server info"), + Implementation::new("http-subscription-server", "1.0.0") + ); + + client.cancel().await?; + server_ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn modern_http_stream_close_without_result_is_abrupt() -> anyhow::Result<()> { + let (url, server_ct, _, subscription_started, _) = spawn_server(ServerEnding::Abrupt).await; + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(url), + ); + let client = ClientConfig::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await?; + let mut subscription = client + .listen(SubscriptionFilter::builder().tools_list_changed().build()) + .await?; + assert!(subscription.next().await?.is_some()); + subscription_started.notified().await; + + server_ct.cancel(); + assert!( + tokio::time::timeout(Duration::from_secs(5), subscription.next()) + .await?? + .is_none() + ); + assert!(matches!(subscription.end(), Some(SubscriptionEnd::Abrupt))); + + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn modern_http_lifecycle_stays_sessionless_for_older_application_version() +-> anyhow::Result<()> { + let (url, server_ct, _, _, get_requests) = spawn_server(ServerEnding::ClientCancellation).await; + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(url), + ); + let client = ClientConfig::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2025_11_25], + }, + ) + .await?; + + client.list_tools(None).await?; + assert_eq!(get_requests.load(Ordering::Relaxed), 0); + + client.cancel().await?; + server_ct.cancel(); + Ok(()) +} + +#[tokio::test] +async fn modern_http_get_and_delete_are_method_not_allowed_in_legacy_session_mode() { + let (url, server_ct, _, _, _) = spawn_server(ServerEnding::ClientCancellation).await; + let client = reqwest::Client::new(); + + for method in [reqwest::Method::GET, reqwest::Method::DELETE] { + let response = client + .request(method, &url) + .header("Accept", "text/event-stream") + .header("MCP-Protocol-Version", "2026-07-28") + .send() + .await + .expect("request"); + assert_eq!(response.status(), reqwest::StatusCode::METHOD_NOT_ALLOWED); + assert_eq!( + response + .headers() + .get(reqwest::header::ALLOW) + .and_then(|value| value.to_str().ok()), + Some("POST") + ); + } + + server_ct.cancel(); +} diff --git a/crates/rmcp/tests/test_task.rs b/crates/rmcp/tests/test_task.rs index ca0f4af..0178bab 100644 --- a/crates/rmcp/tests/test_task.rs +++ b/crates/rmcp/tests/test_task.rs @@ -1,87 +1,423 @@ -use std::{any::Any, time::Duration}; +//! End-to-end tests for the MCP Tasks extension (SEP-2663, +//! `io.modelcontextprotocol/tasks`). +#![cfg(all(feature = "server", feature = "client", not(feature = "local")))] use rmcp::{ - model::TaskStatusNotificationParam, - task_manager::{ - OperationDescriptor, OperationMessage, OperationProcessor, OperationResultTransport, - }, + ErrorData as McpError, ServerHandler, ServiceExt, + handler::server::{router::tool::ToolRouter, wrapper::Parameters}, + model::*, + service::{RequestContext, RoleServer}, + task_manager::{TaskExit, TaskManager, TaskOptions}, + tool, tool_router, }; use serde_json::json; -struct DummyTransport { - id: String, - value: u32, +#[derive(Debug, serde::Deserialize, rmcp::schemars::JsonSchema)] +struct SumArgs { + a: i32, + b: i32, } -impl OperationResultTransport for DummyTransport { - fn operation_id(&self) -> &String { - &self.id +#[derive(Clone)] +struct TaskServer { + tool_router: ToolRouter, + tasks: TaskManager, +} + +#[tool_router] +impl TaskServer { + fn new() -> Self { + Self { + tool_router: Self::tool_router(), + tasks: TaskManager::new(), + } } - fn as_any(&self) -> &dyn Any { - self + #[tool(description = "Sum two numbers")] + async fn sum( + &self, + Parameters(SumArgs { a, b }): Parameters, + ) -> Result { + Ok(CallToolResult::success(vec![ContentBlock::text( + (a + b).to_string(), + )])) } } +impl ServerHandler for TaskServer { + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + let client_supports_tasks = context + .client_capabilities() + .is_some_and(|caps| caps.supports_tasks()); + + if request.name == "sum" && client_supports_tasks { + let args: SumArgs = serde_json::from_value(serde_json::Value::Object( + request.arguments.clone().unwrap_or_default(), + )) + .map_err(|e| McpError::invalid_params(e.to_string(), None))?; + let task = self + .tasks + .spawn(TaskOptions::new().with_poll_interval_ms(10), move |ctx| { + Box::pin(async move { + tokio::select! { + _ = ctx.cancelled() => { + Err(TaskExit::Cancelled) + } + _ = tokio::time::sleep(std::time::Duration::from_millis(50)) => { + Ok(CallToolResult::success(vec![ContentBlock::text( + (args.a + args.b).to_string(), + )])) + } + } + }) + }); + return Ok(CallToolResponse::Task(CreateTaskResult::new(task))); + } + + let tcc = rmcp::handler::server::tool::ToolCallContext::new(self, request, context); + self.tool_router.call(tcc).await + } + + async fn get_task( + &self, + request: GetTaskParams, + _context: RequestContext, + ) -> Result { + Ok(GetTaskResult::new(self.tasks.get_task(&request.task_id)?)) + } + + async fn update_task( + &self, + request: UpdateTaskParams, + _context: RequestContext, + ) -> Result<(), McpError> { + self.tasks + .update_task(&request.task_id, request.input_responses) + } + + async fn cancel_task( + &self, + request: CancelTaskParams, + _context: RequestContext, + ) -> Result<(), McpError> { + self.tasks.cancel_task(&request.task_id) + } + + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tasks() + .build(), + ) + } +} + +fn tasks_client_info() -> ClientConfig { + ClientConfig::new( + ClientCapabilities::builder().enable_tasks().build(), + Implementation::from_build_env(), + ) +} + #[tokio::test] -async fn executes_enqueued_future() { - let mut processor = OperationProcessor::new(); - let descriptor = OperationDescriptor::new("op1", "dummy"); - let future = Box::pin(async { - tokio::time::sleep(Duration::from_millis(10)).await; - Ok(Box::new(DummyTransport { - id: "op1".to_string(), - value: 42, - }) as Box) +async fn task_lifecycle_create_poll_complete() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server = tokio::spawn(async move { + let service = TaskServer::new().serve(server_transport).await?; + service.waiting().await?; + anyhow::Ok(()) }); - processor - .submit_operation(OperationMessage::new(descriptor, future)) - .expect("submit operation"); + let client = tasks_client_info().serve(client_transport).await.unwrap(); - tokio::time::sleep(Duration::from_millis(30)).await; - let results = processor.peek_completed(); - assert_eq!(results.len(), 1); - let payload = results[0] - .result - .as_ref() - .unwrap() - .as_any() - .downcast_ref::() + // Server materializes a task because we declared the extension. + let response = client + .call_tool_once( + CallToolRequestParams::new("sum") + .with_arguments(serde_json::from_value(json!({"a": 40, "b": 2})).unwrap()), + ) + .await .unwrap(); - assert_eq!(payload.value, 42); + let create = match response { + CallToolResponse::Task(create) => create, + other => panic!("expected CreateTaskResult, got {other:?}"), + }; + assert_eq!(create.result_type, ResultType::TASK); + let task_id = create.task.task_id.clone(); + + // Poll until terminal. + let final_task = loop { + tokio::time::sleep(std::time::Duration::from_millis( + create.task.poll_interval_ms.unwrap_or(10), + )) + .await; + let info = client + .peer() + .get_task(GetTaskParams::new(task_id.clone())) + .await + .unwrap(); + if info.task.status().is_terminal() { + break info.task; + } + }; + + match final_task.payload { + TaskPayload::Completed { result } => { + let result: CallToolResult = + serde_json::from_value(serde_json::Value::Object(result)).unwrap(); + let text = result.content[0].as_text().unwrap(); + assert_eq!(text.text, "42"); + } + other => panic!("expected completed task, got {other:?}"), + } + + client.cancel().await.unwrap(); + server.abort(); } #[tokio::test] -async fn rejects_duplicate_operation_ids() { - let mut processor = OperationProcessor::new(); - let descriptor = OperationDescriptor::new("dup", "dummy"); - let future = Box::pin(async { - Ok(Box::new(DummyTransport { - id: "dup".to_string(), - value: 1, - }) as Box) - }); - processor - .submit_operation(OperationMessage::new(descriptor, future)) - .expect("first submit"); - - let descriptor_dup = OperationDescriptor::new("dup", "dummy"); - let future_dup = Box::pin(async { - Ok(Box::new(DummyTransport { - id: "dup".to_string(), - value: 2, - }) as Box) +async fn task_cancel_acknowledged() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server = tokio::spawn(async move { + let service = TaskServer::new().serve(server_transport).await?; + service.waiting().await?; + anyhow::Ok(()) }); - let err = processor - .submit_operation(OperationMessage::new(descriptor_dup, future_dup)) - .expect_err("duplicate should fail"); - assert!(format!("{err}").contains("already running")); + let client = tasks_client_info().serve(client_transport).await.unwrap(); + + let response = client + .call_tool_once( + CallToolRequestParams::new("sum") + .with_arguments(serde_json::from_value(json!({"a": 1, "b": 1})).unwrap()), + ) + .await + .unwrap(); + let create = match response { + CallToolResponse::Task(create) => create, + other => panic!("expected CreateTaskResult, got {other:?}"), + }; + + client + .peer() + .cancel_task(CancelTaskParams::new(create.task.task_id.clone())) + .await + .unwrap(); + + // Cancellation is cooperative (SEP-2663): the ack is immediate, and the + // operation settles the terminal status; poll until it does. + let mut final_status = None; + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + let info = client + .peer() + .get_task(GetTaskParams::new(create.task.task_id.clone())) + .await + .unwrap(); + if info.task.status().is_terminal() { + final_status = Some(info.task.status()); + break; + } + } + assert_eq!(final_status, Some(TaskStatus::Cancelled)); + + client.cancel().await.unwrap(); + server.abort(); +} + +#[tokio::test] +async fn no_task_without_extension_capability() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server = tokio::spawn(async move { + let service = TaskServer::new().serve(server_transport).await?; + service.waiting().await?; + anyhow::Ok(()) + }); + + // Plain client: no tasks extension declared. + let client = ().serve(client_transport).await.unwrap(); + let result = client + .call_tool( + CallToolRequestParams::new("sum") + .with_arguments(serde_json::from_value(json!({"a": 2, "b": 3})).unwrap()), + ) + .await + .unwrap(); + let text = result.content[0].as_text().unwrap(); + assert_eq!(text.text, "5"); + + client.cancel().await.unwrap(); + server.abort(); +} + +/// A misbehaving handler that materializes a task without checking the +/// client's capabilities. The SDK dispatch must catch this and reject with +/// -32021 rather than sending a task handle the client cannot parse. +#[derive(Clone)] +struct AlwaysTaskServer { + tasks: TaskManager, +} + +impl ServerHandler for AlwaysTaskServer { + async fn call_tool( + &self, + _request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + let task = self.tasks.spawn(TaskOptions::default(), |_ctx| { + Box::pin(async { Ok(CallToolResult::success(vec![ContentBlock::text("late")])) }) + }); + Ok(CallToolResponse::Task(CreateTaskResult::new(task))) + } + + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tasks() + .build(), + ) + } +} + +#[tokio::test] +async fn dispatch_rejects_task_result_for_non_declaring_client() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server = tokio::spawn(async move { + let service = AlwaysTaskServer { + tasks: TaskManager::new(), + } + .serve(server_transport) + .await?; + service.waiting().await?; + anyhow::Ok(()) + }); + + // Plain client: no tasks extension declared. The handler tries to return + // a CreateTaskResult anyway; dispatch must reject with -32021. + let client = ().serve(client_transport).await.unwrap(); + let err = client + .call_tool(CallToolRequestParams::new("anything")) + .await + .unwrap_err(); + match err { + rmcp::ServiceError::McpError(e) => { + assert_eq!(e.code, ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY); + } + other => panic!("expected McpError, got {other:?}"), + } + + client.cancel().await.unwrap(); + server.abort(); +} + +#[tokio::test] +async fn tasks_methods_without_capability_return_missing_capability_error() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server = tokio::spawn(async move { + let service = TaskServer::new().serve(server_transport).await?; + service.waiting().await?; + anyhow::Ok(()) + }); + + // Plain client: no tasks extension declared. tasks/* must be rejected + // with -32021 Missing Required Client Capability (SEP-2663), not -32601. + let client = ().serve(client_transport).await.unwrap(); + let err = client + .peer() + .get_task(GetTaskParams::new("whatever")) + .await + .unwrap_err(); + match err { + rmcp::ServiceError::McpError(e) => { + assert_eq!(e.code, ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY); + let data = e.data.expect("error data should be present"); + assert!( + data["requiredCapabilities"]["extensions"] + .as_object() + .is_some_and(|ext| ext.contains_key("io.modelcontextprotocol/tasks")), + "error data should name the tasks extension: {data}" + ); + } + other => panic!("expected McpError, got {other:?}"), + } + + let err = client + .peer() + .cancel_task(CancelTaskParams::new("whatever")) + .await + .unwrap_err(); + match err { + rmcp::ServiceError::McpError(e) => { + assert_eq!(e.code, ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY); + } + other => panic!("expected McpError, got {other:?}"), + } + + client.cancel().await.unwrap(); + server.abort(); +} + +#[tokio::test] +async fn unknown_task_id_returns_invalid_params() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server = tokio::spawn(async move { + let service = TaskServer::new().serve(server_transport).await?; + service.waiting().await?; + anyhow::Ok(()) + }); + + let client = tasks_client_info().serve(client_transport).await.unwrap(); + let err = client + .peer() + .get_task(GetTaskParams::new("no-such-task")) + .await + .unwrap_err(); + match err { + rmcp::ServiceError::McpError(e) => { + // SEP-2663: unknown taskId is -32602 Invalid params. + assert_eq!(e.code, ErrorCode::INVALID_PARAMS); + } + other => panic!("expected McpError, got {other:?}"), + } + + client.cancel().await.unwrap(); + server.abort(); +} + +#[tokio::test] +async fn legacy_task_param_is_ignored() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server = tokio::spawn(async move { + let service = TaskServer::new().serve(server_transport).await?; + service.waiting().await?; + anyhow::Ok(()) + }); + + // Plain client sending a legacy 2025-11-25 `task: {...}` param: it must + // be silently ignored and the call answered synchronously (SEP-2663). + let client = ().serve(client_transport).await.unwrap(); + let params: CallToolRequestParams = serde_json::from_value(json!({ + "name": "sum", + "arguments": {"a": 2, "b": 3}, + "task": {"ttl": 60000} + })) + .expect("legacy task param must not break deserialization"); + let result = client.call_tool(params).await.unwrap(); + let text = result.content[0].as_text().unwrap(); + assert_eq!(text.text, "5"); + + client.cancel().await.unwrap(); + server.abort(); } #[test] -fn task_status_notification_param_preserves_meta() { +fn task_status_notification_params_preserve_meta() { let raw = json!({ "_meta": { "traceId": "trace-1" @@ -90,17 +426,17 @@ fn task_status_notification_param_preserves_meta() { "status": "working", "createdAt": "2026-06-24T00:00:00Z", "lastUpdatedAt": "2026-06-24T00:00:01Z", - "ttl": null + "ttlMs": null }); - let params: TaskStatusNotificationParam = serde_json::from_value(raw).unwrap(); + let params: TaskStatusNotificationParams = serde_json::from_value(raw).unwrap(); - assert_eq!(params.task.task_id, "task-1"); - assert_eq!(params.task_id, "task-1"); + assert_eq!(params.task.task.task_id, "task-1"); + assert_eq!(params.status(), TaskStatus::Working); assert_eq!(params.meta.as_ref().unwrap().0["traceId"], json!("trace-1")); let serialized = serde_json::to_value(¶ms).unwrap(); - assert_eq!(serialized["_meta"]["traceId"], json!("trace-1")); assert_eq!(serialized["taskId"], json!("task-1")); + assert_eq!(serialized["ttlMs"], serde_json::Value::Null); } diff --git a/crates/rmcp/tests/test_task_support_validation.rs b/crates/rmcp/tests/test_task_support_validation.rs deleted file mode 100644 index 41d03f0..0000000 --- a/crates/rmcp/tests/test_task_support_validation.rs +++ /dev/null @@ -1,251 +0,0 @@ -#![cfg(not(feature = "local"))] -//! Tests for task support validation in tool calls. -//! -//! Verifies that the server correctly validates `execution.taskSupport` settings -//! per the MCP specification: -//! - `Required`: MUST be invoked as a task, returns -32601 otherwise -//! - `Forbidden`: MUST NOT be invoked as a task, returns error otherwise -//! - `Optional`: MAY be invoked either way -#![cfg(feature = "client")] - -use rmcp::{ - ClientHandler, ServerHandler, ServiceError, ServiceExt, - handler::server::router::tool::ToolRouter, - model::{CallToolRequestParams, ClientInfo, ErrorCode, TaskMetadata}, - tool, tool_handler, tool_router, -}; - -/// Server with tools having different task support modes. -#[derive(Debug, Clone)] -pub struct TaskSupportTestServer { - #[expect(dead_code, reason = "tool_handler macro accesses this router field")] - tool_router: ToolRouter, -} - -impl Default for TaskSupportTestServer { - fn default() -> Self { - Self::new() - } -} - -impl TaskSupportTestServer { - pub fn new() -> Self { - Self { - tool_router: Self::tool_router(), - } - } -} - -#[tool_router] -impl TaskSupportTestServer { - #[tool( - description = "Tool that requires task-based invocation", - execution(task_support = "required") - )] - async fn required_task_tool(&self) -> String { - "required task executed".to_string() - } - - #[tool( - description = "Tool that forbids task-based invocation", - execution(task_support = "forbidden") - )] - async fn forbidden_task_tool(&self) -> String { - "forbidden task executed".to_string() - } - - #[tool( - description = "Tool that optionally supports task-based invocation", - execution(task_support = "optional") - )] - async fn optional_task_tool(&self) -> String { - "optional task executed".to_string() - } -} - -#[tool_handler] -impl ServerHandler for TaskSupportTestServer {} - -#[derive(Debug, Clone, Default)] -struct DummyClientHandler {} - -impl ClientHandler for DummyClientHandler { - fn get_info(&self) -> ClientInfo { - ClientInfo::default() - } -} - -/// Helper to create a task object for tool calls -fn make_task() -> TaskMetadata { - TaskMetadata::new() -} - -#[tokio::test] -async fn test_required_task_tool_without_task_returns_method_not_found() -> anyhow::Result<()> { - let (server_transport, client_transport) = tokio::io::duplex(4096); - - let server = TaskSupportTestServer::new(); - let server_handle = tokio::spawn(async move { - server.serve(server_transport).await?.waiting().await?; - anyhow::Ok(()) - }); - - let client_handler = DummyClientHandler::default(); - let client = client_handler.serve(client_transport).await?; - - // Call the task-required tool without a task - should fail with -32601 - let result = client - .call_tool(CallToolRequestParams::new("required_task_tool")) - .await; - - // Should be an error with code -32601 (METHOD_NOT_FOUND) - assert!( - result.is_err(), - "Expected error for required task tool without task" - ); - let error = result.unwrap_err(); - - // Check the error data contains the expected code - match error { - ServiceError::McpError(error_data) => { - assert_eq!( - error_data.code, - ErrorCode::METHOD_NOT_FOUND, - "Expected METHOD_NOT_FOUND error code (-32601)" - ); - assert!( - error_data - .message - .contains("requires task-based invocation"), - "Error message should indicate task-based invocation is required, got: {}", - error_data.message - ); - } - _ => panic!("Expected McpError variant, got: {:?}", error), - } - - client.cancel().await?; - server_handle.await??; - Ok(()) -} - -#[tokio::test] -async fn test_forbidden_task_tool_with_task_returns_error() -> anyhow::Result<()> { - let (server_transport, client_transport) = tokio::io::duplex(4096); - - let server = TaskSupportTestServer::new(); - let server_handle = tokio::spawn(async move { - server.serve(server_transport).await?.waiting().await?; - anyhow::Ok(()) - }); - - let client_handler = DummyClientHandler::default(); - let client = client_handler.serve(client_transport).await?; - - // Call the forbidden task tool WITH a task - should fail - let result = client - .call_tool(CallToolRequestParams::new("forbidden_task_tool").with_task(make_task())) - .await; - - // Should be an error with code INVALID_PARAMS - assert!( - result.is_err(), - "Expected error for forbidden task tool with task" - ); - let error = result.unwrap_err(); - - // Check the error data contains the expected code - match error { - ServiceError::McpError(error_data) => { - assert_eq!( - error_data.code, - ErrorCode::INVALID_PARAMS, - "Expected INVALID_PARAMS error code" - ); - assert!( - error_data - .message - .contains("does not support task-based invocation"), - "Error message should indicate task-based invocation is not supported, got: {}", - error_data.message - ); - } - _ => panic!("Expected McpError variant, got: {:?}", error), - } - - client.cancel().await?; - server_handle.await??; - Ok(()) -} - -#[tokio::test] -async fn test_forbidden_task_tool_without_task_succeeds() -> anyhow::Result<()> { - let (server_transport, client_transport) = tokio::io::duplex(4096); - - let server = TaskSupportTestServer::new(); - let server_handle = tokio::spawn(async move { - server.serve(server_transport).await?.waiting().await?; - anyhow::Ok(()) - }); - - let client_handler = DummyClientHandler::default(); - let client = client_handler.serve(client_transport).await?; - - // Call the forbidden task tool WITHOUT a task - should succeed - let result = client - .call_tool(CallToolRequestParams::new("forbidden_task_tool")) - .await; - - assert!( - result.is_ok(), - "Forbidden task tool without task should succeed" - ); - let result = result.unwrap(); - let text = result - .content - .first() - .and_then(|c| c.as_text()) - .map(|t| t.text.as_str()) - .unwrap_or(""); - assert_eq!(text, "forbidden task executed"); - - client.cancel().await?; - server_handle.await??; - Ok(()) -} - -#[tokio::test] -async fn test_optional_task_tool_without_task_succeeds() -> anyhow::Result<()> { - let (server_transport, client_transport) = tokio::io::duplex(4096); - - let server = TaskSupportTestServer::new(); - let server_handle = tokio::spawn(async move { - server.serve(server_transport).await?.waiting().await?; - anyhow::Ok(()) - }); - - let client_handler = DummyClientHandler::default(); - let client = client_handler.serve(client_transport).await?; - - // Call the optional task tool WITHOUT a task - should succeed - let result = client - .call_tool(CallToolRequestParams::new("optional_task_tool")) - .await; - - assert!( - result.is_ok(), - "Optional task tool without task should succeed" - ); - let result = result.unwrap(); - let text = result - .content - .first() - .and_then(|c| c.as_text()) - .map(|t| t.text.as_str()) - .unwrap_or(""); - assert_eq!(text, "optional task executed"); - - client.cancel().await?; - server_handle.await??; - Ok(()) -} diff --git a/crates/rmcp/tests/test_tool_builder_methods.rs b/crates/rmcp/tests/test_tool_builder_methods.rs index 8be7e5c..cdb1473 100644 --- a/crates/rmcp/tests/test_tool_builder_methods.rs +++ b/crates/rmcp/tests/test_tool_builder_methods.rs @@ -22,10 +22,8 @@ fn test_with_output_schema() { assert!(tool.output_schema.is_some()); - // Verify the schema contains expected fields - let schema_str = serde_json::to_string(tool.output_schema.as_ref().unwrap()).unwrap(); - assert!(schema_str.contains("greeting")); - assert!(schema_str.contains("is_adult")); + let schema = tool.output_schema.as_ref().unwrap(); + assert_eq!(schema.get("type"), Some(&serde_json::json!("object"))); } #[test] @@ -57,7 +55,51 @@ fn test_chained_builder_methods() { assert!(input_schema_str.contains("name")); assert!(input_schema_str.contains("age")); - let output_schema_str = serde_json::to_string(tool.output_schema.as_ref().unwrap()).unwrap(); - assert!(output_schema_str.contains("greeting")); - assert!(output_schema_str.contains("is_adult")); + let output_schema = tool.output_schema.as_ref().unwrap(); + assert_eq!( + output_schema.get("type"), + Some(&serde_json::json!("object")) + ); +} + +#[test] +fn test_with_output_schema_primitive() { + let tool = Tool::new("test", "Test tool", JsonObject::new()).with_output_schema::(); + + assert!(tool.output_schema.is_some()); + + let schema = tool.output_schema.as_ref().unwrap(); + assert_eq!(schema.get("type"), Some(&serde_json::json!("integer"))); + // title should be stripped from output schema + assert!(schema.get("title").is_none()); +} + +#[test] +fn test_with_output_schema_array() { + let tool = + Tool::new("test", "Test tool", JsonObject::new()).with_output_schema::>(); + + assert!(tool.output_schema.is_some()); + + let schema_str = serde_json::to_string(tool.output_schema.as_ref().unwrap()).unwrap(); + assert!(schema_str.contains("\"type\":\"array\"")); + assert!(schema_str.contains("items")); + // title should be stripped from output schema + assert!(!schema_str.contains("title")); +} + +#[test] +fn test_with_output_schema_option() { + let tool = + Tool::new("test", "Test tool", JsonObject::new()).with_output_schema::>(); + + assert!(tool.output_schema.is_some()); + + let schema_str = serde_json::to_string(tool.output_schema.as_ref().unwrap()).unwrap(); + // Option generates a composition schema (anyOf/oneOf/type array with null) + assert!( + schema_str.contains("anyOf") || schema_str.contains("oneOf") || schema_str.contains("null"), + "Expected composition schema for Option, got: {schema_str}" + ); + assert!(!schema_str.contains("title")); } diff --git a/crates/rmcp/tests/test_tool_disable_notification.rs b/crates/rmcp/tests/test_tool_disable_notification.rs index b30a58e..14a8757 100644 --- a/crates/rmcp/tests/test_tool_disable_notification.rs +++ b/crates/rmcp/tests/test_tool_disable_notification.rs @@ -9,7 +9,7 @@ use std::sync::{ use rmcp::{ ClientHandler, RoleClient, RoleServer, ServerHandler, ServiceExt, handler::server::{router::tool::ToolRoute, tool::ToolCallContext}, - model::{CallToolResult, ServerCapabilities, ServerInfo, Tool}, + model::{CallToolResponse, CallToolResult, ServerCapabilities, ServerConfig, Tool}, service::{MaybeSendFuture, NotificationContext}, }; use tokio::sync::{Notify, RwLock}; @@ -26,11 +26,11 @@ impl TestToolServer { let mut tool_router = rmcp::handler::server::router::tool::ToolRouter::::new(); tool_router.add_route(ToolRoute::new_dyn( Tool::new("tool_a", "Tool A", Arc::new(Default::default())), - |_ctx| Box::pin(async { Ok(CallToolResult::default()) }), + |_ctx| Box::pin(async { Ok(CallToolResult::default().into()) }), )); tool_router.add_route(ToolRoute::new_dyn( Tool::new("tool_b", "Tool B", Arc::new(Default::default())), - |_ctx| Box::pin(async { Ok(CallToolResult::default()) }), + |_ctx| Box::pin(async { Ok(CallToolResult::default().into()) }), )); Self { router: Arc::new(RwLock::new(tool_router)), @@ -41,15 +41,15 @@ impl TestToolServer { } impl ServerHandler for TestToolServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) } async fn call_tool( &self, request: rmcp::model::CallToolRequestParams, context: rmcp::service::RequestContext, - ) -> Result { + ) -> Result { let router = self.router.read().await; let tcc = ToolCallContext::new(self, request, context); router.call(tcc).await diff --git a/crates/rmcp/tests/test_tool_macros.rs b/crates/rmcp/tests/test_tool_macros.rs index 89846fa..8642496 100644 --- a/crates/rmcp/tests/test_tool_macros.rs +++ b/crates/rmcp/tests/test_tool_macros.rs @@ -10,7 +10,7 @@ use std::sync::Arc; use rmcp::{ ClientHandler, ServerHandler, ServiceExt, handler::server::{router::tool::ToolRouter, wrapper::Parameters}, - model::{CallToolRequestParams, ClientInfo, ServerCapabilities, ServerInfo}, + model::{CallToolRequestParams, ClientConfig, ServerCapabilities, ServerConfig}, tool, tool_handler, tool_router, }; use schemars::JsonSchema; @@ -287,8 +287,8 @@ fn test_optional_field_schema_generation_via_macro() { struct DummyClientHandler {} impl ClientHandler for DummyClientHandler { - fn get_info(&self) -> ClientInfo { - ClientInfo::default() + fn get_info(&self) -> ClientConfig { + ClientConfig::default() } } @@ -397,8 +397,8 @@ fn test_minimal_server_get_info_auto_generated() { "prompts should not be auto-enabled" ); assert!( - info.capabilities.tasks.is_none(), - "tasks should not be auto-enabled" + !info.capabilities.supports_tasks(), + "tasks extension should not be auto-enabled" ); assert!( !info.server_info.name.is_empty(), @@ -549,8 +549,8 @@ impl ManualInfoServer { #[tool_handler] impl ServerHandler for ManualInfoServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_resources() @@ -573,3 +573,47 @@ fn test_manual_get_info_not_overridden() { "manual resources should be preserved" ); } + +/// Server whose tools come from a `macro_rules!` helper wrapping the whole annotated impl. +#[derive(Debug, Clone)] +struct MacroGeneratedServer; + +macro_rules! define_tools { + ($($name:ident => $description:literal),* $(,)?) => { + #[tool_router] + impl MacroGeneratedServer { + $( + #[tool(description = $description)] + async fn $name(&self) -> String { + stringify!($name).to_owned() + } + )* + } + }; +} + +define_tools!(probe => "what a capability would own"); + +#[test] +fn test_macro_rules_around_the_impl_registers_tools() { + let tools = MacroGeneratedServer::tool_router().list_all(); + + assert_eq!(tools.len(), 1); + assert_eq!(tools[0].name, "probe"); + assert_eq!( + tools[0].description.as_deref(), + Some("what a capability would own") + ); +} + +/// Server that opts in to a router with no tools. +#[derive(Debug, Clone)] +struct EmptyRouterServer; + +#[tool_router(allow_empty)] +impl EmptyRouterServer {} + +#[test] +fn test_allow_empty_builds_a_router_without_tools() { + assert!(EmptyRouterServer::tool_router().list_all().is_empty()); +} diff --git a/crates/rmcp/tests/test_tool_result_meta.rs b/crates/rmcp/tests/test_tool_result_meta.rs index a1d3d3a..37585ac 100644 --- a/crates/rmcp/tests/test_tool_result_meta.rs +++ b/crates/rmcp/tests/test_tool_result_meta.rs @@ -1,14 +1,15 @@ -use rmcp::model::{CallToolResult, ContentBlock, Meta}; +use rmcp::model::{CallToolResult, ContentBlock, MetaObject}; use serde_json::{Value, json}; #[test] fn serialize_tool_result_with_meta() { let content = vec![ContentBlock::text("ok")]; - let mut meta = Meta::new(); + let mut meta = MetaObject::new(); meta.insert("foo".to_string(), json!("bar")); let result = CallToolResult::success(content).with_meta(Some(meta)); let v = serde_json::to_value(&result).unwrap(); let expected = json!({ + "resultType": "complete", "content": [{"type":"text","text":"ok"}], "isError": false, "_meta": {"foo":"bar"} diff --git a/crates/rmcp/tests/test_tool_routers.rs b/crates/rmcp/tests/test_tool_routers.rs index d2bbe86..12a5d40 100644 --- a/crates/rmcp/tests/test_tool_routers.rs +++ b/crates/rmcp/tests/test_tool_routers.rs @@ -21,12 +21,6 @@ struct Request { fields: HashMap, } -#[derive(Debug, schemars::JsonSchema, serde::Deserialize, serde::Serialize)] -struct Sum { - a: i32, - b: i32, -} - #[rmcp::tool_router(router = test_router_1)] impl TestHandler { #[rmcp::tool] diff --git a/crates/rmcp/tests/test_trace_context.rs b/crates/rmcp/tests/test_trace_context.rs index 50214b7..15354a4 100644 --- a/crates/rmcp/tests/test_trace_context.rs +++ b/crates/rmcp/tests/test_trace_context.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use rmcp::{ RoleServer, ServerHandler, ServiceExt, - model::{ClientRequest, CustomRequest, CustomResult, Meta}, + model::{ClientRequest, CustomRequest, CustomResult, RequestMetaObject}, service::{PeerRequestOptions, RequestContext}, }; use serde_json::json; @@ -17,7 +17,7 @@ const BAGGAGE: &str = "userId=alice,region=us-east-1"; /// Records the `_meta` it receives on the incoming request so the test can assert passthrough. struct TraceCapturingServer { receive_signal: Arc, - seen: Arc>>, + seen: Arc>>, } impl ServerHandler for TraceCapturingServer { @@ -56,7 +56,7 @@ async fn trace_context_meta_survives_round_trip() -> anyhow::Result<()> { let client = ().serve(client_transport).await?; // Client attaches trace context to the outgoing request's `_meta`. - let mut meta = Meta::new(); + let mut meta = RequestMetaObject::new(); meta.set_traceparent(TRACEPARENT); meta.set_tracestate(TRACESTATE); meta.set_baggage(BAGGAGE); diff --git a/crates/rmcp/tests/test_unix_socket_transport.rs b/crates/rmcp/tests/test_unix_socket_transport.rs index 4c4ad52..134bb1e 100644 --- a/crates/rmcp/tests/test_unix_socket_transport.rs +++ b/crates/rmcp/tests/test_unix_socket_transport.rs @@ -13,9 +13,12 @@ use http::{HeaderName, HeaderValue}; use hyper_util::rt::TokioIo; use rmcp::{ ServiceExt, + model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId}, transport::{ StreamableHttpClientTransport, UnixSocketHttpClient, - streamable_http_client::StreamableHttpClientTransportConfig, + streamable_http_client::{ + StreamableHttpClient, StreamableHttpClientTransportConfig, StreamableHttpError, + }, }, }; use serde_json::json; @@ -35,10 +38,10 @@ async fn mcp_handler( let mut headers_map = HashMap::new(); for (name, value) in headers.iter() { let name_str = name.as_str(); - if name_str.starts_with("x-") || name_str == "host" { - if let Ok(v) = value.to_str() { - headers_map.insert(name_str.to_string(), v.to_string()); - } + if (name_str.starts_with("x-") || name_str == "host") + && let Ok(v) = value.to_str() + { + headers_map.insert(name_str.to_string(), v.to_string()); } } @@ -46,46 +49,46 @@ async fn mcp_handler( stored.extend(headers_map); drop(stored); - if let Ok(json_body) = serde_json::from_slice::(&body) { - if let Some(method) = json_body.get("method").and_then(|m| m.as_str()) { - if method == "initialize" { - state.initialize_called.notify_one(); - let response = json!({ - "jsonrpc": "2.0", - "id": json_body.get("id"), - "result": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "serverInfo": { - "name": "test-unix-server", - "version": "1.0.0" - } + if let Ok(json_body) = serde_json::from_slice::(&body) + && let Some(method) = json_body.get("method").and_then(|m| m.as_str()) + { + if method == "initialize" { + state.initialize_called.notify_one(); + let response = json!({ + "jsonrpc": "2.0", + "id": json_body.get("id"), + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "serverInfo": { + "name": "test-unix-server", + "version": "1.0.0" } - }); - return ( - StatusCode::OK, - [ - (http::header::CONTENT_TYPE, "application/json"), - ( - http::HeaderName::from_static("mcp-session-id"), - "unix-test-session", - ), - ], - response.to_string(), - ); - } else if method == "notifications/initialized" { - return ( - StatusCode::ACCEPTED, - [ - (http::header::CONTENT_TYPE, "application/json"), - ( - http::HeaderName::from_static("mcp-session-id"), - "unix-test-session", - ), - ], - String::new(), - ); - } + } + }); + return ( + StatusCode::OK, + [ + (http::header::CONTENT_TYPE, "application/json"), + ( + http::HeaderName::from_static("mcp-session-id"), + "unix-test-session", + ), + ], + response.to_string(), + ); + } else if method == "notifications/initialized" { + return ( + StatusCode::ACCEPTED, + [ + (http::header::CONTENT_TYPE, "application/json"), + ( + http::HeaderName::from_static("mcp-session-id"), + "unix-test-session", + ), + ], + String::new(), + ); } } @@ -296,3 +299,56 @@ async fn test_unix_socket_convenience_constructor() -> anyhow::Result<()> { Ok(()) } + +/// A request POST that gets 200 + `application/json` + `{}` must error. +/// Mapping that body to Accepted would hang the caller waiting on SSE. +#[tokio::test] +async fn test_unix_socket_request_malformed_json_is_unexpected() -> anyhow::Result<()> { + let dir = std::env::temp_dir().join(format!("rmcp-test-bad-json-{}", std::process::id())); + std::fs::create_dir_all(&dir)?; + let socket_path = dir.join("mcp.sock"); + let _ = std::fs::remove_file(&socket_path); + + async fn bad_json_handler() -> impl IntoResponse { + ( + StatusCode::OK, + [(http::header::CONTENT_TYPE, "application/json")], + "{}", + ) + } + + let app = Router::new().route("/mcp", post(bad_json_handler)); + let listener = tokio::net::UnixListener::bind(&socket_path)?; + let server_handle = spawn_unix_server(listener, app); + + let socket_str = socket_path.to_str().unwrap(); + let uri = "http://localhost/mcp"; + let client = UnixSocketHttpClient::new(socket_str, uri); + let result = client + .post_message( + Arc::from(uri), + ClientJsonRpcMessage::request( + ClientRequest::PingRequest(PingRequest::default()), + RequestId::Number(1), + ), + None, + None, + HashMap::new(), + ) + .await; + + server_handle.abort(); + let _ = std::fs::remove_file(&socket_path); + let _ = std::fs::remove_dir(&dir); + + match result { + Err(StreamableHttpError::UnexpectedServerResponse(ref msg)) => { + assert!( + msg.contains("{}"), + "expected body preview in error, got: {msg}" + ); + } + other => panic!("expected UnexpectedServerResponse, got: {other:?}"), + } + Ok(()) +} diff --git a/crates/rmcp/tests/test_with_js.rs b/crates/rmcp/tests/test_with_js.rs index 0dbd93f..8992f81 100644 --- a/crates/rmcp/tests/test_with_js.rs +++ b/crates/rmcp/tests/test_with_js.rs @@ -17,6 +17,23 @@ use common::calculator::Calculator; const STREAMABLE_HTTP_BIND_ADDRESS: &str = "127.0.0.1:8001"; const STREAMABLE_HTTP_JS_BIND_ADDRESS: &str = "127.0.0.1:8002"; +// These tests run concurrently, while npm mutates their shared node_modules directory. +static JS_DEPENDENCIES_INSTALLED: tokio::sync::OnceCell<()> = tokio::sync::OnceCell::const_new(); + +async fn install_js_dependencies() -> anyhow::Result<()> { + JS_DEPENDENCIES_INSTALLED + .get_or_try_init(|| async { + let status = tokio::process::Command::new("npm") + .arg("install") + .current_dir("tests/test_with_js") + .status() + .await?; + anyhow::ensure!(status.success(), "npm install failed with {status}"); + Ok(()) + }) + .await?; + Ok(()) +} #[tokio::test] async fn test_with_js_stdio_server() -> anyhow::Result<()> { @@ -27,12 +44,7 @@ async fn test_with_js_stdio_server() -> anyhow::Result<()> { ) .with(tracing_subscriber::fmt::layer()) .try_init(); - tokio::process::Command::new("npm") - .arg("install") - .current_dir("tests/test_with_js") - .spawn()? - .wait() - .await?; + install_js_dependencies().await?; let transport = TokioChildProcess::new(tokio::process::Command::new("node").configure(|cmd| { cmd.arg("tests/test_with_js/server.js"); @@ -57,12 +69,7 @@ async fn test_with_js_streamable_http_client() -> anyhow::Result<()> { ) .with(tracing_subscriber::fmt::layer()) .try_init(); - tokio::process::Command::new("npm") - .arg("install") - .current_dir("tests/test_with_js") - .spawn()? - .wait() - .await?; + install_js_dependencies().await?; let ct = CancellationToken::new(); let service: StreamableHttpService = @@ -104,12 +111,7 @@ async fn test_with_js_streamable_http_server() -> anyhow::Result<()> { ) .with(tracing_subscriber::fmt::layer()) .try_init(); - tokio::process::Command::new("npm") - .arg("install") - .current_dir("tests/test_with_js") - .spawn()? - .wait() - .await?; + install_js_dependencies().await?; let transport = StreamableHttpClientTransport::from_uri(format!( "http://{STREAMABLE_HTTP_JS_BIND_ADDRESS}/mcp" diff --git a/docs/OAUTH_SUPPORT.md b/docs/OAUTH_SUPPORT.md index 16809a4..ffd3bf2 100644 --- a/docs/OAUTH_SUPPORT.md +++ b/docs/OAUTH_SUPPORT.md @@ -1,6 +1,6 @@ # Model Context Protocol OAuth Authorization -This document describes the OAuth 2.1 authorization implementation for Model Context Protocol (MCP), following the [MCP 2025-11-25 Authorization Specification](https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization/). +This document describes the OAuth 2.1 authorization implementation for Model Context Protocol (MCP), following the [MCP Authorization Specification](https://modelcontextprotocol.io/specification/2026-07-28/basic/authorization/). ## Features @@ -15,6 +15,7 @@ This document describes the OAuth 2.1 authorization implementation for Model Con - Automatic token refresh - Authorized HTTP Client implementation - Injectable OAuth HTTP client for custom network environments +- Opt-in EMA/XAA refresh-token and ID-JAG exchanges for registered public and confidential clients ## Usage Guide @@ -61,7 +62,9 @@ have been obtained. If OAuth requests must run outside reqwest, implement `OAuthHttpClient` and use `OAuthState::new_with_oauth_http_client`. The SDK passes each OAuth request to your implementation with the raw HTTP request, a suggested timeout, and an -`OAuthHttpRedirectPolicy`. +`OAuthHttpRedirectPolicy`. `OAuthHttpClientFuture` returns +`OAuthHttpClientError`, so implementations can propagate their native error +types with `?` without flattening their source chains into strings. ```rust ignore use std::sync::Arc; @@ -108,16 +111,76 @@ Use this path when OAuth traffic must go through a browser fetch API, a remote execution environment, a company gateway, a test fake, or any other non-reqwest transport. -### 3. Start authorization with OAuthState +#### Inspect discovery provenance directly -The `OAuthState` state machine manages the full authorization lifecycle. When no -scopes are provided, the SDK automatically selects scopes from the server's -WWW-Authenticate header, Protected Resource Metadata, or AS metadata. +Most applications can use `OAuthState` without calling metadata discovery +directly. When using `AuthorizationManager`, `resolve_metadata()` returns both +the metadata and how it was obtained. A client that supports the 2025-03-26 +default-endpoint fallback can continue with synthesized metadata, while a +client that requires server-published metadata should reject that result: ```rust ignore -// start authorization - pass empty scopes to let the SDK auto-select +use rmcp::transport::auth::{AuthorizationManager, AuthorizationMetadataSource}; + +async fn configure_metadata( + manager: &mut AuthorizationManager, + allow_legacy_endpoint_fallback: bool, +) -> anyhow::Result<()> { + let resolution = manager.resolve_metadata().await?; + + if resolution.source == AuthorizationMetadataSource::LegacyEndpointFallback { + if !allow_legacy_endpoint_fallback { + anyhow::bail!("the server did not publish OAuth metadata"); + } + + tracing::warn!( + "the server did not publish OAuth metadata; using the 2025-03-26 fallback endpoints" + ); + } + + manager.set_metadata(resolution.metadata); + Ok(()) +} +``` + +`ProtectedResourceMetadata` and `AuthorizationServerMetadata` indicate +server-published metadata, so clients can proceed with the returned metadata. +`LegacyEndpointFallback` indicates endpoints synthesized for compatibility +with the 2025-03-26 MCP specification. Clients should proceed only when they +intentionally support that legacy behavior; clients using discovery as an +OAuth capability check should treat it as unsupported. + +Applications using `OAuthState` do not need to handle these sources directly: +the state machine resolves metadata internally and retains the legacy fallback. +Low-level `AuthorizationManager` users can use +`AuthorizationMetadataSource::is_discovered()` when they only need to +distinguish server-published metadata from synthesized metadata. + +### 3. Start authorization with OAuthState + +The `OAuthState` state machine manages the full authorization lifecycle. +`start_authorization` accepts an `AuthorizationRequest` describing the client +identity material you have available, and selects a client registration +mechanism following the [spec's priority order](https://modelcontextprotocol.io/specification/2026-07-28/basic/authorization/client-registration): + +1. **Pre-registered client information** (`with_preregistered_client`), when + the client already holds a `client_id` issued out of band +2. **Client ID Metadata Documents** (SEP-991, `with_client_metadata_url`), when + the authorization server advertises `client_id_metadata_document_supported` +3. **Dynamic Client Registration**, as a fallback when the authorization server + advertises a `registration_endpoint` + +When no scopes are provided, the SDK automatically selects scopes from the +server's WWW-Authenticate header, Protected Resource Metadata, or AS metadata. + +```rust ignore +use rmcp::transport::auth::AuthorizationRequest; + +// start authorization - pass no scopes to let the SDK auto-select oauth_state - .start_authorization(&[], MCP_REDIRECT_URI, Some("My MCP Client")) + .start_authorization( + AuthorizationRequest::new(MCP_REDIRECT_URI).with_client_name("My MCP Client"), + ) .await .context("Failed to start authorization")?; ``` @@ -126,7 +189,40 @@ If you know the scopes you need, you can still pass them explicitly: ```rust ignore oauth_state - .start_authorization(&["mcp", "profile"], MCP_REDIRECT_URI, Some("My MCP Client")) + .start_authorization( + AuthorizationRequest::new(MCP_REDIRECT_URI) + .with_scopes(["mcp", "profile"]) + .with_client_name("My MCP Client"), + ) + .await + .context("Failed to start authorization")?; +``` + +If the client hosts a Client ID Metadata Document (SEP-991), pass its URL; the +SDK uses it when the server supports CIMD and falls back to dynamic +registration otherwise: + +```rust ignore +oauth_state + .start_authorization( + AuthorizationRequest::new(MCP_REDIRECT_URI) + .with_client_name("My MCP Client") + .with_client_metadata_url("https://example.com/client-metadata.json"), + ) + .await + .context("Failed to start authorization")?; +``` + +If the client was registered with the authorization server out of band, provide +the pre-registered credentials; they take priority over every other mechanism: + +```rust ignore +oauth_state + .start_authorization( + AuthorizationRequest::new(MCP_REDIRECT_URI) + .with_preregistered_client("my-client-id") + .with_client_secret("my-client-secret"), + ) .await .context("Failed to start authorization")?; ``` @@ -157,10 +253,30 @@ let transport = StreamableHttpClientTransport::with_client( ); // create client and connect to MCP server -let client_service = ClientInfo::default(); +let client_service = ClientConfig::default(); let client = client_service.serve(transport).await?; ``` +If initialization reports that authorization is required, return to the +application's authorization flow: + +```rust ignore +let client = match client_service.serve(transport).await { + Ok(client) => client, + Err(error) if error.is_authorization_required() => { + // Prompt the user and start the application's authorization flow again. + return Err(error.into()); + } + Err(error) => return Err(error.into()), +}; +``` + +The predicate covers both missing or expired local OAuth authorization and an +HTTP 401 challenge from the MCP server. Other failures, including transient +token-refresh errors and insufficient scope, return `false`. The original error +is preserved for logging or more detailed handling; the SDK does not start an +authorization flow automatically. + ### 6. Handle scope upgrades If a server returns 403 with `insufficient_scope`, you can request a scope @@ -179,9 +295,134 @@ match oauth_state.request_scope_upgrade("admin:write", MCP_REDIRECT_URI).await { } ``` +## Enterprise-managed authorization (EMA/XAA) + +The example requires the `rmcp` features `auth-enterprise-managed`, `client`, +`reqwest` (TLS), and `transport-streamable-http-client-reqwest`, plus `oauth2` +version 5. Call the async function from a Tokio runtime. + +The exchange profile has these requirements and limits: + +- Each authorization server has its own approved client registration and explicit + `EmaClientAuthentication`: `None`, `ClientSecretBasic`, `ClientSecretPost`, or + `JwtAssertion`. The SDK does not select methods from metadata or fall back to a + different method after a failure. +- Input is an enterprise IdP refresh token. The requested MCP resource must match + the ID-JAG's sole `resource` claim; scope may be omitted or narrowed. +- RAR (`authorization_details`) and DPoP are not supported. Nonempty authorization + details are rejected at both exchange stages. +- Redemption consumes the SDK's ID-JAG handle and does not retry automatically. + This is an SDK safety choice, not a protocol requirement that ID-JAGs be single-use. + +The [ID-JAG draft recommends confidential clients](https://datatracker.ietf.org/doc/html/draft-ietf-oauth-identity-assertion-authz-grant-04#section-9.1). +The example below uses confidential clients with `client_secret_basic` at both +servers. Use `None` only where that server permits a public client registration. +Discovery, server approval, SSO, credential storage, and reauthentication remain +the application's responsibility. Client-side ID-JAG checks validate structure and +bindings, not signatures; the resource authorization server verifies signatures. + +```rust no_run +use oauth2::{ClientSecret, RefreshToken}; +use rmcp::{ + ServiceExt, + model::ClientConfig, + transport::{ + StreamableHttpClientTransport, + auth::{ + default_oauth_http_client, + enterprise::{EmaAuthorizationServer, EmaClientAuthentication, EmaExchangeRequest}, + }, + streamable_http_client::StreamableHttpClientTransportConfig, + }, +}; + +async fn connect( + refresh: &RefreshToken, + idp_client_secret: ClientSecret, + resource_client_secret: ClientSecret, +) -> Result<(), Box> { + let resource = "https://mcp.example/mcp"; + let http = default_oauth_http_client()?; + let idp = EmaAuthorizationServer::new( + "https://idp.example", "https://idp.example/token", "idp-client", + ) + .with_client_authentication(EmaClientAuthentication::ClientSecretBasic(idp_client_secret)); + let resource_as = EmaAuthorizationServer::new( + "https://as.example", "https://as.example/token", "mcp-client", + ) + .with_client_authentication(EmaClientAuthentication::ClientSecretBasic(resource_client_secret)); + let token = EmaExchangeRequest::new(idp, resource_as, resource, refresh) + .with_scopes(["files.read"]) + .exchange(&http, &http) + .await?; + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(resource) + .auth_header(token.access_token.secret()), + ); + let client = ClientConfig::default().serve(transport).await?; + client.list_tools(Default::default()).await?; + client.cancel().await?; + Ok(()) +} +``` + +`auth_header` takes the token without a `Bearer ` prefix. Use it only with the +approved resource and never log it. This transport uses a fixed token; obtain a +new token and reconnect when it expires or is rejected. + +For a registration using JWT client authentication, implement +`EmaClientAssertionProvider` with your application's signer and configure +`JwtAssertion` on that server. The provider example below also uses `async-trait` +version 0.1; `AppSigner` represents your application's existing signing service. + +```rust ignore +use std::sync::Arc; +use rmcp::transport::auth::enterprise::{ + EmaAuthorizationServer, EmaClientAssertion, EmaClientAssertionProvider, + EmaClientAuthentication, +}; + +struct AppAssertionProvider { + signer: AppSigner, +} + +#[async_trait::async_trait] +impl EmaClientAssertionProvider for AppAssertionProvider { + async fn create_assertion( + &self, + server: &EmaAuthorizationServer, + ) -> Result> { + // Sign a new assertion for this registration and approved server. + let jwt = self.signer.sign_client_assertion( + &server.client_id, &server.issuer, &server.token_endpoint, + ).await?; + Ok(EmaClientAssertion::new(jwt)) + } +} + +let resource_as = EmaAuthorizationServer::new( + "https://as.example", "https://as.example/token", "mcp-client", +) +.with_client_authentication(EmaClientAuthentication::JwtAssertion(Arc::new( + AppAssertionProvider { signer }, +))); +``` + +The SDK calls the provider before each token request, including delayed +`EmaIdJag::exchange` redemption. Sign a fresh, short-lived assertion with a unique +`jti`, the registered client ID in `iss` and `sub`, and the server's approved +audience in `aud`. Your signer owns the keys and algorithm; the client assertion +is separate from the ID-JAG grant. Signing and HTTP share a 30-second deadline. + +The factory honors per-request redirect policy with the SDK's default reqwest +settings. For custom proxy, CA, or remote-execution policy, implement +`OAuthHttpClient`; use separate adapters for the IdP and resource AS when their +network policies differ. + ## Complete Examples -- **Client**: [`examples/clients/src/auth/oauth_client.rs`](../examples/clients/src/auth/oauth_client.rs) +- **Authorization Code client**: [`examples/clients/src/auth/oauth_client.rs`](../examples/clients/src/auth/oauth_client.rs) +- **Client Credentials client**: [`examples/clients/src/auth/client_credentials.rs`](../examples/clients/src/auth/client_credentials.rs) - **Server**: [`examples/servers/src/complex_auth_streamhttp.rs`](../examples/servers/src/complex_auth_streamhttp.rs) ### Running the Examples @@ -192,6 +433,10 @@ cargo run -p mcp-server-examples --example servers_complex_auth_streamhttp # Run the OAuth client (in another terminal) cargo run -p mcp-client-examples --example clients_oauth_client + +# Run the Client Credentials client +cargo run -p mcp-client-examples --example clients_client_credentials -- \ + ``` ## Authorization Flow Description @@ -231,7 +476,7 @@ If you encounter authorization issues, check the following: ## References -- [MCP Authorization Specification (2025-11-25)](https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization/) +- [MCP Authorization Specification](https://modelcontextprotocol.io/specification/2026-07-28/basic/authorization/) - [OAuth 2.1 Specification Draft](https://oauth.net/2.1/) - [RFC 8414: OAuth 2.0 Authorization Server Metadata](https://datatracker.ietf.org/doc/html/rfc8414) - [RFC 7591: OAuth 2.0 Dynamic Client Registration Protocol](https://datatracker.ietf.org/doc/html/rfc7591) diff --git a/docs/readme/README.zh-cn.md b/docs/readme/README.zh-cn.md deleted file mode 100644 index 8439384..0000000 --- a/docs/readme/README.zh-cn.md +++ /dev/null @@ -1,990 +0,0 @@ -

- -# RMCP -[![Crates.io Version](https://img.shields.io/crates/v/rmcp)](https://crates.io/crates/rmcp) -[![docs.rs](https://img.shields.io/docsrs/rmcp)](https://docs.rs/rmcp/latest/rmcp) -[![CI](https://github.com/modelcontextprotocol/rust-sdk/actions/workflows/ci.yml/badge.svg)](https://github.com/modelcontextprotocol/rust-sdk/actions/workflows/ci.yml) -[![License](https://img.shields.io/crates/l/rmcp)](../../LICENSE) - -一个基于 tokio 异步运行时的官方 Rust Model Context Protocol SDK 实现。 - -> **迁移到 1.x？** 请参阅 [迁移指南](https://github.com/modelcontextprotocol/rust-sdk/discussions/716) 了解破坏性变更和升级说明。 - -本仓库包含以下 crate： - -- [rmcp](../../crates/rmcp)：实现 RMCP 协议的核心库 - 详见 [rmcp](../../crates/rmcp/README.md) -- [rmcp-macros](../../crates/rmcp-macros)：用于生成 RMCP 工具实现的过程宏库 - 详见 [rmcp-macros](../../crates/rmcp-macros/README.md) - -完整的 MCP 规范请参阅 [modelcontextprotocol.io](https://modelcontextprotocol.io/specification/2025-11-25)。 - -## 目录 - -- [使用](#使用) -- [工具](#工具) -- [资源](#资源) -- [提示词](#提示词) -- [采样](#采样) -- [根目录](#根目录) -- [日志](#日志) -- [补全](#补全) -- [通知](#通知) -- [订阅](#订阅) -- [任务](#任务长时间运行的工具调用) -- [示例](#示例) -- [OAuth 支持](#oauth-支持) -- [相关资源](#相关资源) -- [相关项目](#相关项目) -- [开发](#开发) - -## 使用 - -### 导入 - -```toml -rmcp = { version = "0.16.0", features = ["server"] } -## 或使用最新开发版本 -rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", branch = "main" } -``` -### 第三方依赖 - -基本依赖： -- [tokio](https://github.com/tokio-rs/tokio) -- [serde](https://github.com/serde-rs/serde) -JSON Schema 生成 (version 2020-12)： -- [schemars](https://github.com/GREsau/schemars) - -### 构建客户端 - -
-启动客户端 - -```rust, ignore -use rmcp::{ServiceExt, transport::{TokioChildProcess, ConfigureCommandExt}}; -use tokio::process::Command; - -#[tokio::main] -async fn main() -> Result<(), Box> { - let client = ().serve(TokioChildProcess::new(Command::new("npx").configure(|cmd| { - cmd.arg("-y").arg("@modelcontextprotocol/server-everything"); - }))?).await?; - Ok(()) -} -``` -
- -### 构建服务端 - -
-构建传输层 - -```rust, ignore -use tokio::io::{stdin, stdout}; -let transport = (stdin(), stdout()); -``` - -
- -
-构建服务 - -你可以通过 [`ServerHandler`](../../crates/rmcp/src/handler/server.rs) 或 [`ClientHandler`](../../crates/rmcp/src/handler/client.rs) 轻松构建服务。 - -```rust, ignore -let service = common::counter::Counter::new(); -``` -
- -
-启动服务端 - -```rust, ignore -// 此调用将完成初始化过程 -let server = service.serve(transport).await?; -``` -
- -
-与服务端交互 - -服务端初始化完成后，你可以发送请求或通知： - -```rust, ignore -// 请求 -let roots = server.list_roots().await?; - -// 或发送通知 -server.notify_cancelled(...).await?; -``` -
- -
-等待服务停止 - -```rust, ignore -let quit_reason = server.waiting().await?; -// 或将其取消 -let quit_reason = server.cancel().await?; -``` -
- ---- - -## 工具 - -工具允许服务端向客户端暴露可调用的函数。每个工具都有名称、描述和参数的 JSON Schema。客户端通过 `list_tools` 发现工具，通过 `call_tool` 调用工具。 - -**MCP 规范：** [Tools](https://modelcontextprotocol.io/specification/2025-11-25/server/tools) - -### 服务端 - -`#[tool]`、`#[tool_router]` 和 `#[tool_handler]` 宏负责所有连接工作。对于纯工具服务端，可以使用 `#[tool_router(server_handler)]` 来省略单独的 `ServerHandler` 实现： - -```rust,ignore -use rmcp::{handler::server::wrapper::Parameters, schemars, tool, tool_router, ServiceExt, transport::stdio}; - -#[derive(Debug, serde::Deserialize, schemars::JsonSchema)] -struct AddParams { - a: i32, - b: i32, -} - -#[derive(Clone)] -struct Calculator; - -#[tool_router(server_handler)] -impl Calculator { - #[tool(description = "Add two numbers")] - fn add(&self, Parameters(AddParams { a, b }): Parameters) -> String { - (a + b).to_string() - } -} - -#[tokio::main] -async fn main() -> anyhow::Result<()> { - let service = Calculator.serve(stdio()).await?; - service.waiting().await?; - Ok(()) -} -``` - -当需要自定义服务端元数据或多种能力（工具 + 提示词）时，使用显式的 `#[tool_handler]`： - -```rust,ignore -use rmcp::{handler::server::wrapper::Parameters, schemars, tool, tool_router, tool_handler, ServerHandler, ServiceExt}; - -#[derive(Debug, serde::Deserialize, schemars::JsonSchema)] -struct AddParams { - a: i32, - b: i32, -} - -#[derive(Clone)] -struct Calculator; - -#[tool_router] -impl Calculator { - #[tool(description = "Add two numbers")] - fn add(&self, Parameters(AddParams { a, b }): Parameters) -> String { - (a + b).to_string() - } -} - -#[tool_handler(name = "calculator", version = "1.0.0", instructions = "A simple calculator")] -impl ServerHandler for Calculator {} -``` - -完整的宏文档请参阅 [`crates/rmcp-macros`](../../crates/rmcp-macros/README.md)。 - -### 客户端 - -```rust,ignore -use rmcp::model::CallToolRequestParams; - -// 列出所有工具 -let tools = client.list_all_tools().await?; - -// 按名称调用工具 -let result = client.call_tool(CallToolRequestParams::new("add")).await?; -``` - -**示例：** [`examples/servers/src/common/calculator.rs`](../../examples/servers/src/common/calculator.rs)（服务端），[`examples/servers/src/calculator_stdio.rs`](../../examples/servers/src/calculator_stdio.rs)（stdio 运行器） - ---- - -## 资源 - -资源允许服务端向客户端暴露数据（文件、数据库记录、API 响应）供其读取。每个资源通过 URI 标识，返回文本或二进制（base64 编码）内容。资源模板允许服务端声明带有动态参数的 URI 模式。 - -**MCP 规范：** [Resources](https://modelcontextprotocol.io/specification/2025-11-25/server/resources) - -### 服务端 - -在 `ServerHandler` trait 上实现 `list_resources()`、`read_resource()`，以及可选的 `list_resource_templates()`。在 `get_info()` 中启用资源能力。 - -```rust -use rmcp::{ - ErrorData as McpError, RoleServer, ServerHandler, ServiceExt, - model::*, - service::RequestContext, - transport::stdio, -}; -use serde_json::json; - -#[derive(Clone)] -struct MyServer; - -impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( - ServerCapabilities::builder() - .enable_resources() - .build(), - ) - } - - async fn list_resources( - &self, - _request: Option, - _context: RequestContext, - ) -> Result { - Ok(ListResourcesResult { - resources: vec![ - Resource::new("file:///config.json", "config"), - Resource::new("memo://insights", "insights"), - ], - next_cursor: None, - meta: None, - }) - } - - async fn read_resource( - &self, - request: ReadResourceRequestParams, - _context: RequestContext, - ) -> Result { - match request.uri.as_str() { - "file:///config.json" => Ok(ReadResourceResult::new(vec![ - ResourceContents::text(r#"{"key": "value"}"#, &request.uri), - ])), - "memo://insights" => Ok(ReadResourceResult::new(vec![ - ResourceContents::text("Analysis results...", &request.uri), - ])), - _ => Err(McpError::resource_not_found( - "resource_not_found", - Some(json!({ "uri": request.uri })), - )), - } - } - - async fn list_resource_templates( - &self, - _request: Option, - _context: RequestContext, - ) -> Result { - Ok(ListResourceTemplatesResult { - resource_templates: vec![], - next_cursor: None, - meta: None, - }) - } -} -``` - -### 客户端 - -```rust -use rmcp::model::{ReadResourceRequestParams}; - -// 列出所有资源（自动处理分页） -let resources = client.list_all_resources().await?; - -// 通过 URI 读取特定资源 -let result = client.read_resource( - ReadResourceRequestParams::new("file:///config.json"), -).await?; - -// 列出资源模板 -let templates = client.list_all_resource_templates().await?; -``` - -### 通知 - -服务端可以在资源列表变更或特定资源更新时通知客户端： - -```rust -// 通知资源列表已变更（客户端应重新获取） -context.peer.notify_resource_list_changed().await?; - -// 通知特定资源已更新 -context.peer.notify_resource_updated( - ResourceUpdatedNotificationParam::new("file:///config.json"), -).await?; -``` - -客户端通过 `ClientHandler` 处理这些通知： - -```rust -impl ClientHandler for MyClient { - async fn on_resource_list_changed( - &self, - _context: NotificationContext, - ) { - // 重新获取资源列表 - } - - async fn on_resource_updated( - &self, - params: ResourceUpdatedNotificationParam, - _context: NotificationContext, - ) { - // 重新读取 params.uri 对应的资源 - } -} -``` - -**示例：** [`examples/servers/src/common/counter.rs`](../../examples/servers/src/common/counter.rs)（服务端），[`examples/clients/src/everything_stdio.rs`](../../examples/clients/src/everything_stdio.rs)（客户端） - ---- - -## 提示词 - -提示词是服务端向客户端暴露的可复用消息模板。它们接受类型化参数并返回对话消息。`#[prompt]` 宏自动处理参数验证和路由。 - -**MCP 规范：** [Prompts](https://modelcontextprotocol.io/specification/2025-11-25/server/prompts) - -### 服务端 - -使用 `#[prompt_router]`、`#[prompt]` 和 `#[prompt_handler]` 宏以声明式方式定义提示词。参数定义为派生 `JsonSchema` 的结构体。 - -```rust -use rmcp::{ - ErrorData as McpError, RoleServer, ServerHandler, ServiceExt, - handler::server::{router::prompt::PromptRouter, wrapper::Parameters}, - model::*, - prompt, prompt_handler, prompt_router, - schemars::JsonSchema, - service::RequestContext, - transport::stdio, -}; -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Serialize, Deserialize, JsonSchema)] -pub struct CodeReviewArgs { - #[schemars(description = "Programming language of the code")] - pub language: String, - #[schemars(description = "Focus areas for the review")] - pub focus_areas: Option>, -} - -#[derive(Clone)] -pub struct MyServer { - prompt_router: PromptRouter, -} - -#[prompt_router] -impl MyServer { - fn new() -> Self { - Self { prompt_router: Self::prompt_router() } - } - - /// 无参数的简单提示词 - #[prompt(name = "greeting", description = "A simple greeting")] - async fn greeting(&self) -> Vec { - vec![PromptMessage::new_text( - Role::User, - "Hello! How can you help me today?", - )] - } - - /// 带类型化参数的提示词 - #[prompt(name = "code_review", description = "Review code in a given language")] - async fn code_review( - &self, - Parameters(args): Parameters, - ) -> Result { - let focus = args.focus_areas - .unwrap_or_else(|| vec!["correctness".into()]); - - Ok(GetPromptResult::new(vec![ - PromptMessage::new_text( - Role::User, - format!("Review my {} code. Focus on: {}", args.language, focus.join(", ")), - ), - ]) - .with_description(format!("Code review for {}", args.language))) - } -} - -#[prompt_handler] -impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_prompts().build()) - } -} -``` - -提示词函数支持以下返回类型： -- `Vec` -- 简单消息列表 -- `GetPromptResult` -- 带可选描述的消息 -- `Result` -- 以上任一类型，附带错误处理 - -### 客户端 - -```rust -use rmcp::model::GetPromptRequestParams; - -// 列出所有提示词 -let prompts = client.list_all_prompts().await?; - -// 带参数获取提示词 -let result = client.get_prompt(GetPromptRequestParams { - meta: None, - name: "code_review".into(), - arguments: Some(rmcp::object!({ - "language": "Rust", - "focus_areas": ["performance", "safety"] - })), -}).await?; -``` - -### 通知 - -```rust -// 服务端：通知可用提示词已变更 -context.peer.notify_prompt_list_changed().await?; -``` - -**示例：** [`examples/servers/src/prompt_stdio.rs`](../../examples/servers/src/prompt_stdio.rs)（服务端），[`examples/clients/src/everything_stdio.rs`](../../examples/clients/src/everything_stdio.rs)（客户端） - ---- - -## 采样 - -采样反转了通常的方向：服务端请求客户端执行 LLM 补全。服务端发送 `create_message` 请求，客户端通过其 LLM 处理并返回结果。 - -**MCP 规范：** [Sampling](https://modelcontextprotocol.io/specification/2025-11-25/client/sampling) - -### 服务端（请求采样） - -通过 `context.peer.create_message()` 访问客户端的采样能力： - -```rust -use rmcp::model::*; - -// 在 ServerHandler 方法内部（例如 call_tool）： -let response = context.peer.create_message( - CreateMessageRequestParams::new( - vec![SamplingMessage::user_text("Explain this error: connection refused")], - 150, - ) - .with_model_preferences( - ModelPreferences::new() - .with_hints(vec![ModelHint::new("claude")]) - .with_cost_priority(0.3) - .with_speed_priority(0.8) - .with_intelligence_priority(0.7), - ) - .with_system_prompt("You are a helpful assistant.") - .with_include_context(ContextInclusion::None) - .with_temperature(0.7), -).await?; - -// 提取响应文本 -let text = response.message.content - .first() - .and_then(|c| c.as_text()) - .map(|t| &t.text); -``` - -### 客户端（处理采样） - -在客户端实现 `ClientHandler::create_message()`。这是你调用实际 LLM 的地方： - -```rust -use rmcp::{ClientHandler, model::*, service::{RequestContext, RoleClient}}; - -#[derive(Clone, Default)] -struct MyClient; - -impl ClientHandler for MyClient { - async fn create_message( - &self, - params: CreateMessageRequestParams, - _context: RequestContext, - ) -> Result { - // 转发到你的 LLM，或返回模拟响应： - let response_text = call_your_llm(¶ms.messages).await; - - Ok(CreateMessageResult::new( - SamplingMessage::assistant_text(response_text), - "my-model".into(), - ) - .with_stop_reason(CreateMessageResult::STOP_REASON_END_TURN)) - } -} -``` - -**示例：** [`examples/servers/src/sampling_stdio.rs`](../../examples/servers/src/sampling_stdio.rs)（服务端），[`examples/clients/src/sampling_stdio.rs`](../../examples/clients/src/sampling_stdio.rs)（客户端） - ---- - -## 根目录 - -根目录告诉服务端客户端正在使用哪些目录或项目。根目录是一个 URI（通常为 `file://`），指向工作区或代码仓库。服务端可以查询根目录以了解在哪里查找文件以及如何限定工作范围。 - -**MCP 规范：** [Roots](https://modelcontextprotocol.io/specification/2025-11-25/client/roots) - -### 服务端 - -向客户端请求根目录列表，并处理变更通知： - -```rust -use rmcp::{ServerHandler, model::*, service::{NotificationContext, RoleServer}}; - -impl ServerHandler for MyServer { - // 向客户端查询根目录 - async fn call_tool( - &self, - request: CallToolRequestParams, - context: RequestContext, - ) -> Result { - let roots = context.peer.list_roots().await?; - // 使用 roots.roots 了解工作区边界 - // ... - } - - // 当客户端的根目录列表变更时调用 - async fn on_roots_list_changed( - &self, - _context: NotificationContext, - ) { - // 重新获取根目录以保持最新 - } -} -``` - -### 客户端 - -客户端声明根目录能力并实现 `list_roots()`： - -```rust -use rmcp::{ClientHandler, model::*}; - -impl ClientHandler for MyClient { - async fn list_roots( - &self, - _context: RequestContext, - ) -> Result { - Ok(ListRootsResult::new(vec![ - Root::new("file:///home/user/project").with_name("My Project"), - ])) - } -} -``` - -客户端在根目录变更时通知服务端： - -```rust -// 添加或移除工作区根目录后： -client.notify_roots_list_changed().await?; -``` - ---- - -## 日志 - -服务端可以向客户端发送结构化日志消息。客户端设置最低严重级别，服务端通过对等通知接口发送消息。 - -**MCP 规范：** [Logging](https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/logging) - -### 服务端 - -启用日志能力，处理客户端的级别变更，并通过对等端发送日志消息： - -```rust -use rmcp::{ServerHandler, model::*, service::RequestContext}; - -impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( - ServerCapabilities::builder() - .enable_logging() - .build(), - ) - } - - // 客户端设置最低日志级别 - async fn set_level( - &self, - request: SetLevelRequestParams, - _context: RequestContext, - ) -> Result<(), ErrorData> { - // 存储 request.level 并据此过滤后续日志消息 - Ok(()) - } -} - -// 在任何可以访问 peer 的处理器中发送日志消息： -context.peer.notify_logging_message( - LoggingMessageNotificationParam::new( - LoggingLevel::Info, - serde_json::json!({ - "message": "Processing completed", - "items_processed": 42 - }), - ) - .with_logger("my-server"), -).await?; -``` - -可用日志级别（从低到高）：`Debug`、`Info`、`Notice`、`Warning`、`Error`、`Critical`、`Alert`、`Emergency`。 - -### 客户端 - -客户端通过 `ClientHandler` 处理传入的日志消息： - -```rust -impl ClientHandler for MyClient { - async fn on_logging_message( - &self, - params: LoggingMessageNotificationParam, - _context: NotificationContext, - ) { - println!("[{}] {}: {}", params.level, - params.logger.unwrap_or_default(), params.data); - } -} -``` - -客户端也可以设置服务端的日志级别： - -```rust -client.set_level(SetLevelRequestParams::new(LoggingLevel::Warning)).await?; -``` - ---- - -## 补全 - -补全为提示词或资源模板参数提供自动补全建议。当用户填写参数时，客户端可以根据已输入的内容向服务端请求建议。 - -**MCP 规范：** [Completions](https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/completion) - -### 服务端 - -启用补全能力并实现 `complete()` 处理器。使用 `request.context` 检查已填写的参数： - -```rust -use rmcp::{ErrorData as McpError, ServerHandler, model::*, service::RequestContext, RoleServer}; - -impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( - ServerCapabilities::builder() - .enable_completions() - .enable_prompts() - .build(), - ) - } - - async fn complete( - &self, - request: CompleteRequestParams, - _context: RequestContext, - ) -> Result { - let values = match &request.r#ref { - Reference::Prompt(prompt_ref) if prompt_ref.name == "sql_query" => { - match request.argument.name.as_str() { - "operation" => vec!["SELECT", "INSERT", "UPDATE", "DELETE"], - "table" => vec!["users", "orders", "products"], - "columns" => { - // 根据已填写的参数调整建议 - if let Some(ctx) = &request.context { - if let Some(op) = ctx.get_argument("operation") { - match op.to_uppercase().as_str() { - "SELECT" | "UPDATE" => { - vec!["id", "name", "email", "created_at"] - } - _ => vec![], - } - } else { vec![] } - } else { vec![] } - } - _ => vec![], - } - } - _ => vec![], - }; - - // 根据用户的部分输入进行过滤 - let filtered: Vec = values.into_iter() - .map(String::from) - .filter(|v| v.to_lowercase().contains(&request.argument.value.to_lowercase())) - .collect(); - - let completion = CompletionInfo::with_pagination(filtered, None, false) - .map_err(|e| McpError::internal_error(e, None))?; - Ok(CompleteResult::new(completion)) - } -} -``` - -### 客户端 - -```rust -use rmcp::model::*; - -let result = client.complete(CompleteRequestParams::new( - Reference::for_prompt("sql_query"), - ArgumentInfo::new("operation", "SEL"), -)).await?; - -// result.completion.values 包含建议，例如 ["SELECT"] -``` - -**示例：** [`examples/servers/src/completion_stdio.rs`](../../examples/servers/src/completion_stdio.rs) - ---- - -## 通知 - -通知是即发即忘的消息——不需要响应。它们涵盖进度更新、取消和生命周期事件。双方都可以发送和接收通知。 - -**MCP 规范：** [Notifications](https://modelcontextprotocol.io/specification/2025-11-25/basic/notifications) - -### 进度通知 - -服务端可以在长时间运行的操作中报告进度： - -```rust -use rmcp::model::*; - -// 在工具处理器内部： -for i in 0..total_items { - process_item(i).await; - - context.peer.notify_progress( - ProgressNotificationParam::new( - ProgressToken(NumberOrString::Number(i as i64)), - i as f64, - ) - .with_total(total_items as f64) - .with_message(format!("Processing item {}/{}", i + 1, total_items)), - ).await?; -} -``` - -### 取消 - -任一方都可以取消正在进行的请求： - -```rust -// 发送取消通知 -context.peer.notify_cancelled(CancelledNotificationParam::new( - Some(the_request_id), - Some("User requested cancellation".into()), -)).await?; -``` - -在 `ServerHandler` 或 `ClientHandler` 中处理取消： - -```rust -impl ServerHandler for MyServer { - async fn on_cancelled( - &self, - params: CancelledNotificationParam, - _context: NotificationContext, - ) { - // 中止 params.request_id 对应的工作 - } -} -``` - -### 初始化通知 - -客户端在握手完成后发送 `initialized` 通知： - -```rust -// 在 serve() 握手过程中由 rmcp 自动发送。 -// 服务端通过以下方式处理： -impl ServerHandler for MyServer { - async fn on_initialized( - &self, - _context: NotificationContext, - ) { - // 服务端已准备好接收请求 - } -} -``` - -### 列表变更通知 - -当可用的工具、提示词或资源发生变更时，通知客户端： - -```rust -context.peer.notify_tool_list_changed().await?; -context.peer.notify_prompt_list_changed().await?; -context.peer.notify_resource_list_changed().await?; -``` - -**示例：** [`examples/servers/src/common/progress_demo.rs`](../../examples/servers/src/common/progress_demo.rs) - ---- - -## 订阅 - -客户端可以订阅特定资源。当订阅的资源发生变更时，服务端发送通知，客户端可以重新读取该资源。 - -**MCP 规范：** [Resources - Subscriptions](https://modelcontextprotocol.io/specification/2025-11-25/server/resources#subscriptions) - -### 服务端 - -在资源能力中启用订阅，并实现 `subscribe()` / `unsubscribe()` 处理器： - -```rust -use rmcp::{ErrorData as McpError, ServerHandler, model::*, service::RequestContext, RoleServer}; -use std::sync::Arc; -use tokio::sync::Mutex; -use std::collections::HashSet; - -#[derive(Clone)] -struct MyServer { - subscriptions: Arc>>, -} - -impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( - ServerCapabilities::builder() - .enable_resources() - .enable_resources_subscribe() - .build(), - ) - } - - async fn subscribe( - &self, - request: SubscribeRequestParams, - _context: RequestContext, - ) -> Result<(), McpError> { - self.subscriptions.lock().await.insert(request.uri); - Ok(()) - } - - async fn unsubscribe( - &self, - request: UnsubscribeRequestParams, - _context: RequestContext, - ) -> Result<(), McpError> { - self.subscriptions.lock().await.remove(&request.uri); - Ok(()) - } -} -``` - -当订阅的资源发生变更时，通知客户端： - -```rust -// 检查资源是否有订阅者，然后通知 -context.peer.notify_resource_updated( - ResourceUpdatedNotificationParam::new("file:///config.json"), -).await?; -``` - -### 客户端 - -```rust -use rmcp::model::*; - -// 订阅资源更新 -client.subscribe(SubscribeRequestParams::new("file:///config.json")).await?; - -// 不再需要时取消订阅 -client.unsubscribe(UnsubscribeRequestParams::new("file:///config.json")).await?; -``` - -在 `ClientHandler` 中处理更新通知： - -```rust -impl ClientHandler for MyClient { - async fn on_resource_updated( - &self, - params: ResourceUpdatedNotificationParam, - _context: NotificationContext, - ) { - // 重新读取 params.uri 对应的资源 - } -} -``` - ---- - -## 任务（长时间运行的工具调用） - -`rmcp` 支持 SEP-1319 中定义的[基于任务的工具调用](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks)流程。为工具添加 `execution(task_support = "required" | "optional")` 注解，并在 `ServerHandler` 实现上添加 `#[task_handler]` —— `enqueue_task`、`tasks/list`、`tasks/get`、`tasks/result` 和 `tasks/cancel` 将在 `OperationProcessor` 之上自动生成。 - -```rust, ignore -#[tool( - description = "Sum two numbers after a 2-second delay", - execution(task_support = "required") -)] -async fn slow_sum(/* ... */) -> Result { /* ... */ } - -#[tool_handler] -#[task_handler] -impl ServerHandler for TaskDemo {} -``` - -完整的端到端示例请参阅 [`servers_task_stdio`](../../examples/servers/src/task_stdio.rs) 及对应的 [`clients_task_stdio`](../../examples/clients/src/task_stdio.rs)。 - ---- - -## 示例 - -查看 [examples](../../examples/README.md)。 - -## OAuth 支持 - -查看 [OAuth 支持](../OAUTH_SUPPORT.md) 了解详情。 - -## 相关资源 - -- [MCP 规范](https://modelcontextprotocol.io/specification/2025-11-25) -- [Schema](https://github.com/modelcontextprotocol/specification/blob/main/schema/2025-11-25/schema.ts) - -## 相关项目 - -### 扩展 `rmcp` - -- [rmcp-actix-web](https://gitlab.com/lx-industries/rmcp-actix-web) - 基于 `actix_web` 的 `rmcp` 后端 -- [rmcp-openapi](https://gitlab.com/lx-industries/rmcp-openapi) - 将 OpenAPI 定义的端点转换为 MCP 工具 - -### 基于 `rmcp` 构建 - -- [goose](https://github.com/block/goose) - 一个超越代码建议的开源、可扩展 AI 智能体 -- [apollo-mcp-server](https://github.com/apollographql/apollo-mcp-server) - 通过 Apollo GraphOS 将 AI 智能体连接到 GraphQL API 的 MCP 服务 -- [rustfs-mcp](https://github.com/rustfs/rustfs/tree/main/crates/mcp) - 为 AI/LLM 集成提供 S3 兼容对象存储操作的高性能 MCP 服务 -- [containerd-mcp-server](https://github.com/jokemanfire/mcp-containerd) - 基于 containerd 实现的 MCP 服务 -- [rmcp-openapi-server](https://gitlab.com/lx-industries/rmcp-openapi/-/tree/main/crates/rmcp-openapi-server) - 将 OpenAPI 定义的端点暴露为 MCP 工具的高性能 MCP 服务 -- [nvim-mcp](https://github.com/linw1995/nvim-mcp) - 与 Neovim 交互的 MCP 服务 -- [terminator](https://github.com/mediar-ai/terminator) - AI 驱动的桌面自动化 MCP 服务，支持跨平台，成功率超过 95% -- [stakpak-agent](https://github.com/stakpak/agent) - 安全加固的 DevOps 终端智能体，支持 MCP over mTLS、流式传输、密钥令牌化和异步任务管理 -- [video-transcriber-mcp-rs](https://github.com/nhatvu148/video-transcriber-mcp-rs) - 使用 whisper.cpp 从 1000+ 平台转录视频的高性能 MCP 服务 -- [NexusCore MCP](https://github.com/sjkim1127/Nexuscore_MCP) - 具有 Frida 集成和隐蔽脱壳功能的高级恶意软件分析与动态检测 MCP 服务 -- [spreadsheet-mcp](https://github.com/PSU3D0/spreadsheet-mcp) - 面向 LLM 智能体的高效 Token 使用的电子表格分析 MCP 服务，支持自动区域检测、重新计算、截图和编辑 -- [hyper-mcp](https://github.com/hyper-mcp-rs/hyper-mcp) - 通过 WebAssembly (WASM) 插件扩展功能的快速、安全的 MCP 服务 -- [rudof-mcp](https://github.com/rudof-project/rudof/tree/master/rudof_mcp) - RDF 验证和数据处理 MCP 服务，支持 ShEx/SHACL 验证、SPARQL 查询和格式转换。支持 stdio 和 Streamable HTTP 传输，具备完整的 MCP 功能（工具、提示词、资源、日志、补全、任务） -- [MCPMate](https://github.com/loocor/MCPMate) - 渐进式 MCP 管理桌面应用：从引导式服务导入开始，逐步扩展到多客户端配置集和 Unify 元工具，让能力暴露、Token 消耗与运行状态更可控，并在效率、成本和可靠性上提供更多选择 - - -## 开发 - -### 贡献指南 - -查看 [docs/CONTRIBUTE.MD](../CONTRIBUTE.MD) 获取贡献提示。 - -### 使用 Dev Container - -如果你想使用 Dev Container，查看 [docs/DEVCONTAINER.md](../DEVCONTAINER.md) 获取开发指南。 diff --git a/examples/clients/Cargo.toml b/examples/clients/Cargo.toml index 416057a..419a176 100644 --- a/examples/clients/Cargo.toml +++ b/examples/clients/Cargo.toml @@ -65,3 +65,7 @@ path = "src/auth/client_credentials.rs" [[example]] name = "clients_task_stdio" path = "src/task_stdio.rs" + +[[example]] +name = "clients_subscriptions_streamhttp" +path = "src/subscriptions_streamhttp.rs" diff --git a/examples/clients/README.md b/examples/clients/README.md index 419cdac..36bdd59 100644 --- a/examples/clients/README.md +++ b/examples/clients/README.md @@ -16,6 +16,13 @@ A client that communicates with a Git-related MCP server using standard input/ou A client that communicates with an MCP server using HTTP streaming transport. - Connects to an MCP server running at `http://localhost:8000` + +### Modern Subscription Client (`subscriptions_streamhttp.rs`) + +Uses modern discovery and `subscriptions/listen`, prints the accepted filter, +and consumes tagged notifications until graceful closure or cancellation. + +- Run with `cargo run -p mcp-client-examples --example clients_subscriptions_streamhttp` - Retrieves server information and list of available tools - Calls a tool named "increment" @@ -50,6 +57,14 @@ A client demonstrating how to authenticate with an MCP server using OAuth. - Establishes an authorized connection to the MCP server using the acquired access token - Demonstrates how to use the authorized connection to retrieve available tools and prompts +### OAuth Client Credentials (`auth/client_credentials.rs`) + +A client demonstrating the OAuth 2.0 Client Credentials flow from SEP-1046. + +- Accepts the server URL, client ID, and client secret as command-line arguments +- Authenticates without an interactive browser or callback server +- Establishes an authorized connection and retrieves the available tools + ### Sampling Standard I/O Client (`sampling_stdio.rs`) @@ -62,12 +77,13 @@ A client demonstrating how to use the sampling tool. ### Task Standard I/O Client (`task_stdio.rs`) -A client that exercises the task lifecycle against `servers_task_stdio` -(per [SEP-1319](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks)). +A client that exercises the SEP-2663 Tasks extension lifecycle against `servers_task_stdio` +([SEP-2663](https://modelcontextprotocol.io/extensions/tasks/overview), `io.modelcontextprotocol/tasks`). - Spawns `servers_task_stdio` as a child process over stdio +- Declares the tasks extension in its client capabilities - Calls `quick_echo` synchronously -- Calls `slow_sum` as a task via `CallToolRequestParams::with_task(...)`, polls `tasks/get` until completion, then fetches the result via `tasks/result` +- Calls `slow_sum`, receives a `CreateTaskResult` (`resultType: "task"`), polls `tasks/get` honoring `pollIntervalMs`, and reads the final `CallToolResult` inlined in the completed task ### Progress Test Client (`progress_client.rs`) @@ -99,6 +115,10 @@ cargo run -p mcp-client-examples --example clients_collection # Run the OAuth client example cargo run -p mcp-client-examples --example clients_oauth_client +# Run the OAuth Client Credentials example +cargo run -p mcp-client-examples --example clients_client_credentials -- \ + + # Run the sampling standard I/O client example cargo run -p mcp-client-examples --example clients_sampling_stdio diff --git a/examples/clients/src/auth/client_credentials.rs b/examples/clients/src/auth/client_credentials.rs index 55aa615..e277b9e 100644 --- a/examples/clients/src/auth/client_credentials.rs +++ b/examples/clients/src/auth/client_credentials.rs @@ -3,7 +3,7 @@ use std::env; use anyhow::{Context, Result}; use rmcp::{ ServiceExt, - model::ClientInfo, + model::ClientConfig, transport::{ StreamableHttpClientTransport, auth::{AuthClient, ClientCredentialsConfig, OAuthState}, @@ -73,7 +73,7 @@ async fn main() -> Result<()> { ); // Connect to MCP server and list tools - let client_service = ClientInfo::default(); + let client_service = ClientConfig::default(); let client = client_service.serve(transport).await?; tracing::info!("Connected to MCP server"); diff --git a/examples/clients/src/auth/oauth_client.rs b/examples/clients/src/auth/oauth_client.rs index ffd3434..8fa18df 100644 --- a/examples/clients/src/auth/oauth_client.rs +++ b/examples/clients/src/auth/oauth_client.rs @@ -8,11 +8,12 @@ use axum::{ routing::get, }; use rmcp::{ - ServiceExt, - model::ClientInfo, + RoleClient, ServiceExt, + model::ClientConfig, + service::RunningService, transport::{ StreamableHttpClientTransport, - auth::{AuthClient, OAuthState}, + auth::{AuthClient, AuthorizationRequest, OAuthState}, streamable_http_client::StreamableHttpClientTransportConfig, }, }; @@ -55,6 +56,109 @@ async fn callback_handler( Html(CALLBACK_HTML.to_string()) } +enum ConnectOutcome { + /// The server accepted the unauthenticated connection. + Connected(RunningService), + /// The server answered 401; authorize with this `WWW-Authenticate` + /// challenge and reconnect. + AuthRequired(String), +} + +/// Attempt the real connection unauthenticated — the reactive discovery +/// trigger (matching the TypeScript and Python SDKs). The server's 401 +/// challenge, not a probe, tells us whether and how to authorize. +async fn try_connect(http_client: reqwest::Client, server_url: &str) -> Result { + let transport = StreamableHttpClientTransport::with_client( + http_client, + StreamableHttpClientTransportConfig::with_uri(server_url), + ); + match ClientConfig::default().serve(transport).await { + Ok(client) => Ok(ConnectOutcome::Connected(client)), + Err(error) => match error.auth_challenge() { + Some(challenge) => Ok(ConnectOutcome::AuthRequired(challenge.to_string())), + None => Err(error.into()), + }, + } +} + +/// Run the browser OAuth flow seeded by the server's challenge, then +/// reconnect with the authorized transport. +async fn authorize_and_connect( + challenge: String, + oauth_http_client: reqwest::Client, + server_url: &str, + client_metadata_url: &str, + code_receiver: oneshot::Receiver, + output: &mut BufWriter, +) -> Result> { + tracing::info!("Server requires authorization: {challenge}"); + + // initialize oauth state machine + let mut oauth_state = OAuthState::new(server_url, Some(oauth_http_client)) + .await + .context("Failed to initialize oauth state machine")?; + // Seed discovery from the server's challenge, and use CIMD (SEP-991) + // with client metadata URL. Passing no scopes lets the SDK auto-select + // from the challenge's scope hint, Protected Resource Metadata, or AS + // metadata. + oauth_state + .start_authorization( + AuthorizationRequest::new(MCP_REDIRECT_URI) + .with_client_name("Test MCP Client") + .with_client_metadata_url(client_metadata_url) + .with_challenge(challenge), + ) + .await + .context("Failed to start authorization")?; + + // Output authorization URL to user + output + .write_all(b"Please open the following URL in your browser to authorize:\n\n") + .await?; + output + .write_all(oauth_state.get_authorization_url().await?.as_bytes()) + .await?; + output + .write_all(b"\n\nWaiting for browser callback, please do not close this window...\n") + .await?; + output.flush().await?; + + // Wait for authorization code + tracing::info!("Waiting for authorization code..."); + let CallbackParams { + code: auth_code, + state: csrf_token, + iss, + } = code_receiver + .await + .context("Failed to get authorization code")?; + tracing::info!("Received authorization code: {}", auth_code); + // Exchange code for access token + tracing::info!("Exchanging authorization code for access token..."); + oauth_state + .handle_callback_with_issuer(&auth_code, &csrf_token, iss.as_deref()) + .await + .context("Failed to handle callback")?; + tracing::info!("Successfully obtained access token"); + + output + .write_all(b"\nAuthorization successful! Access token obtained.\n\n") + .await?; + output.flush().await?; + + // Reconnect with the authorized transport + tracing::info!("Establishing authorized connection to MCP server..."); + let am = oauth_state + .into_authorization_manager() + .ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?; + let auth_client = AuthClient::new(reqwest::Client::default(), am); + let transport = StreamableHttpClientTransport::with_client( + auth_client, + StreamableHttpClientTransportConfig::with_uri(server_url), + ); + Ok(ClientConfig::default().serve(transport).await?) +} + #[tokio::main] async fn main() -> Result<()> { // Initialize logging @@ -123,74 +227,32 @@ async fn main() -> Result<()> { .build() .context("Failed to build OAuth HTTP client")?; - // initialize oauth state machine - let mut oauth_state = OAuthState::new(&server_url, Some(oauth_http_client)) - .await - .context("Failed to initialize oauth state machine")?; - // use CIMD (SEP-991) with client metadata URL. - // passing empty scopes lets the SDK auto-select from the server's - // WWW-Authenticate header, Protected Resource Metadata, or AS metadata. - oauth_state - .start_authorization_with_metadata_url( - &[], - MCP_REDIRECT_URI, - Some("Test MCP Client"), - Some(&client_metadata_url), - ) - .await - .context("Failed to start authorization")?; - - // Output authorization URL to user let mut output = BufWriter::new(tokio::io::stdout()); output.write_all(b"\n=== MCP OAuth Client ===\n\n").await?; - output - .write_all(b"Please open the following URL in your browser to authorize:\n\n") - .await?; - output - .write_all(oauth_state.get_authorization_url().await?.as_bytes()) - .await?; - output - .write_all(b"\n\nWaiting for browser callback, please do not close this window...\n") - .await?; output.flush().await?; - // Wait for authorization code - tracing::info!("Waiting for authorization code..."); - let CallbackParams { - code: auth_code, - state: csrf_token, - iss, - } = code_receiver - .await - .context("Failed to get authorization code")?; - tracing::info!("Received authorization code: {}", auth_code); - // Exchange code for access token - tracing::info!("Exchanging authorization code for access token..."); - oauth_state - .handle_callback_with_issuer(&auth_code, &csrf_token, iss.as_deref()) - .await - .context("Failed to handle callback")?; - tracing::info!("Successfully obtained access token"); - - output - .write_all(b"\nAuthorization successful! Access token obtained.\n\n") - .await?; - output.flush().await?; - - // Create authorized transport, this transport is authorized by the oauth state machine - tracing::info!("Establishing authorized connection to MCP server..."); - let am = oauth_state - .into_authorization_manager() - .ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?; - let client = AuthClient::new(reqwest::Client::default(), am); - let transport = StreamableHttpClientTransport::with_client( - client, - StreamableHttpClientTransportConfig::with_uri(server_url.as_str()), - ); - - // Create client and connect to MCP server - let client_service = ClientInfo::default(); - let client = client_service.serve(transport).await?; + // Reactive discovery: attempt the real connection first. The server's + // 401 challenge — not a probe — tells us whether and how to authorize. + // The transport gets a default client: `oauth_http_client`'s request + // timeout would cut long-lived SSE streams short. + tracing::info!("Attempting connection to MCP server..."); + let client = match try_connect(reqwest::Client::default(), &server_url).await? { + ConnectOutcome::Connected(client) => { + tracing::info!("Server accepted the connection without authorization"); + client + } + ConnectOutcome::AuthRequired(challenge) => { + authorize_and_connect( + challenge, + oauth_http_client, + &server_url, + &client_metadata_url, + code_receiver, + &mut output, + ) + .await? + } + }; tracing::info!("Successfully connected to MCP server"); // Test API requests diff --git a/examples/clients/src/progress_client.rs b/examples/clients/src/progress_client.rs index 89c4873..3990006 100644 --- a/examples/clients/src/progress_client.rs +++ b/examples/clients/src/progress_client.rs @@ -8,7 +8,7 @@ use clap::{Parser, ValueEnum}; use rmcp::{ ClientHandler, ServiceExt, model::{ - CallToolRequestParams, ClientCapabilities, ClientInfo, Implementation, + CallToolRequestParams, ClientCapabilities, ClientConfig, Implementation, ProgressNotificationParam, }, service::{NotificationContext, RoleClient}, @@ -100,10 +100,10 @@ impl ProgressAwareClient { } fn stop_tracking(&self) { - if let Ok(mut tracker_opt) = self.tracker.lock() { - if let Some(tracker) = tracker_opt.take() { - tracker.print_summary(); - } + if let Ok(mut tracker_opt) = self.tracker.lock() + && let Some(tracker) = tracker_opt.take() + { + tracker.print_summary(); } } } @@ -114,15 +114,15 @@ impl ClientHandler for ProgressAwareClient { params: ProgressNotificationParam, _context: NotificationContext, ) { - if let Ok(tracker_opt) = self.tracker.lock() { - if let Some(tracker) = tracker_opt.as_ref() { - tracker.handle_progress(¶ms); - } + if let Ok(tracker_opt) = self.tracker.lock() + && let Some(tracker) = tracker_opt.as_ref() + { + tracker.handle_progress(¶ms); } } - fn get_info(&self) -> ClientInfo { - ClientInfo::new( + fn get_info(&self) -> ClientConfig { + ClientConfig::new( ClientCapabilities::default(), Implementation::new("progress-test-client", "1.0.0"), ) @@ -163,7 +163,10 @@ async fn test_stdio_transport(records: u32) -> Result<()> { // Initialize let server_info = service.peer_info(); if let Some(info) = server_info { - tracing::info!("Connected to server: {:?}", info.server_info.name); + tracing::info!( + "Connected to server: {:?}", + info.server_info.as_ref().map(|server| &server.name) + ); } // List tools @@ -179,10 +182,10 @@ async fn test_stdio_transport(records: u32) -> Result<()> { .call_tool(CallToolRequestParams::new("stream_processor")) .await?; - if let Some(content) = tool_result.content.first() { - if let Some(text) = content.as_text() { - tracing::info!("Processing completed: {}", text.text); - } + if let Some(content) = tool_result.content.first() + && let Some(text) = content.as_text() + { + tracing::info!("Processing completed: {}", text.text); } service.cancel().await?; @@ -214,7 +217,10 @@ async fn test_http_transport(http_url: &str, records: u32) -> Result<()> { // Initialize let server_info = client.peer_info(); if let Some(info) = server_info { - tracing::info!("Connected to server: {:?}", info.server_info.name); + tracing::info!( + "Connected to server: {:?}", + info.server_info.as_ref().map(|server| &server.name) + ); } // List tools @@ -230,10 +236,10 @@ async fn test_http_transport(http_url: &str, records: u32) -> Result<()> { .call_tool(CallToolRequestParams::new("stream_processor")) .await?; - if let Some(content) = tool_result.content.first() { - if let Some(text) = content.as_text() { - tracing::info!("processing completed: {}", text.text); - } + if let Some(content) = tool_result.content.first() + && let Some(text) = content.as_text() + { + tracing::info!("processing completed: {}", text.text); } client.cancel().await?; diff --git a/examples/clients/src/sampling_stdio.rs b/examples/clients/src/sampling_stdio.rs index cc7c5f1..1079586 100644 --- a/examples/clients/src/sampling_stdio.rs +++ b/examples/clients/src/sampling_stdio.rs @@ -1,3 +1,8 @@ +#![expect( + deprecated, + reason = "This example demonstrates the deprecated MCP sampling API" +)] + use anyhow::Result; use rmcp::{ ClientHandler, ServiceExt, diff --git a/examples/clients/src/streamable_http.rs b/examples/clients/src/streamable_http.rs index 0c27fd3..f841a87 100644 --- a/examples/clients/src/streamable_http.rs +++ b/examples/clients/src/streamable_http.rs @@ -1,14 +1,15 @@ use anyhow::Result; use rmcp::{ - ServiceExt, - model::{CallToolRequestParams, ClientCapabilities, ClientInfo, Implementation}, + ClientLifecycleMode, ClientServiceExt, + model::{ + CallToolRequestParams, ClientCapabilities, ClientConfig, Implementation, ProtocolVersion, + }, transport::StreamableHttpClientTransport, }; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; #[tokio::main] async fn main() -> Result<()> { - // Initialize logging tracing_subscriber::registry() .with( tracing_subscriber::EnvFilter::try_from_default_env() @@ -17,27 +18,31 @@ async fn main() -> Result<()> { .with(tracing_subscriber::fmt::layer()) .init(); let transport = StreamableHttpClientTransport::from_uri("http://localhost:8000/mcp"); - let client_info = ClientInfo::new( + let client_info = ClientConfig::new( ClientCapabilities::default(), - Implementation::new("test sse client", "0.0.1"), + Implementation::new("streamable-http-client", "0.0.1"), ); - let client = client_info.serve(transport).await.inspect_err(|e| { - tracing::error!("client error: {:?}", e); - })?; + let client = client_info + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await + .inspect_err(|e| { + tracing::error!("client error: {:?}", e); + })?; - // Initialize let server_info = client.peer_info(); tracing::info!("Connected to server: {server_info:#?}"); - // List tools let tools = client.list_tools(Default::default()).await?; tracing::info!("Available tools: {tools:#?}"); let tool_result = client - .call_tool( - CallToolRequestParams::new("increment") - .with_arguments(serde_json::json!({}).as_object().cloned().unwrap()), - ) + .call_tool(CallToolRequestParams::new("increment")) .await?; tracing::info!("Tool result: {tool_result:#?}"); client.cancel().await?; diff --git a/examples/clients/src/subscriptions_streamhttp.rs b/examples/clients/src/subscriptions_streamhttp.rs new file mode 100644 index 0000000..9c04531 --- /dev/null +++ b/examples/clients/src/subscriptions_streamhttp.rs @@ -0,0 +1,49 @@ +use rmcp::{ + ClientLifecycleMode, ClientServiceExt, + model::{ClientConfig, ProtocolVersion, SubscriptionFilter}, + transport::StreamableHttpClientTransport, +}; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()), + ) + .init(); + + let transport = StreamableHttpClientTransport::from_uri("http://127.0.0.1:8000/mcp"); + let client = ClientConfig::default() + .serve_with_lifecycle( + transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await?; + let mut subscription = client + .listen(SubscriptionFilter::builder().tools_list_changed().build()) + .await?; + + println!("accepted filter: {:?}", subscription.acknowledged()); + loop { + tokio::select! { + result = subscription.next() => { + match result? { + Some(notification) => println!("notification: {notification:?}"), + None => { + println!("subscription ended: {:?}", subscription.end()); + break; + } + } + } + _ = tokio::signal::ctrl_c() => { + subscription.cancel().await?; + break; + } + } + } + + client.cancel().await?; + Ok(()) +} diff --git a/examples/clients/src/task_stdio.rs b/examples/clients/src/task_stdio.rs index ffcc0dc..a946f75 100644 --- a/examples/clients/src/task_stdio.rs +++ b/examples/clients/src/task_stdio.rs @@ -1,18 +1,20 @@ //! Client for the task-demo server in `examples/servers/src/task_stdio.rs`. //! -//! Walks through the task lifecycle (SEP-1319): +//! Walks through the SEP-2663 Tasks extension lifecycle: //! 1. Call a regular tool (`quick_echo`) — synchronous response. -//! 2. Call a task-required tool (`slow_sum`) by attaching `task: {}` to -//! the `tools/call` request. The server returns a `Task` with a `task_id`. -//! 3. Poll `tasks/get` until status becomes `Completed`. -//! 4. Fetch the underlying `CallToolResult` via `tasks/result`. +//! 2. Call `slow_sum` while declaring the `io.modelcontextprotocol/tasks` +//! extension capability. The server decides to materialize a task and +//! returns a `CreateTaskResult` (`resultType: "task"`). +//! 3. Poll `tasks/get` (honoring `pollIntervalMs`) until the task reaches a +//! terminal status; the final `CallToolResult` is inlined in the +//! `completed` task's `result` field. use anyhow::{Result, anyhow}; use rmcp::{ ServiceExt, model::{ - CallToolRequestParams, CallToolResult, ClientRequest, GetTaskParams, GetTaskPayloadParams, - Request, ServerResult, TaskMetadata, TaskStatus, + CallToolRequestParams, CallToolResponse, CallToolResult, ClientCapabilities, GetTaskParams, + TaskPayload, TaskStatus, }, object, transport::{ConfigureCommandExt, TokioChildProcess}, @@ -30,8 +32,14 @@ async fn main() -> Result<()> { .with(tracing_subscriber::fmt::layer()) .init(); + // Declare the tasks extension in our client capabilities (SEP-2663). + let client_info = rmcp::model::ClientConfig::new( + ClientCapabilities::builder().enable_tasks().build(), + rmcp::model::Implementation::from_build_env(), + ); + // Spawn the task-demo server as a child process over stdio. - let client = () + let client = client_info .serve(TokioChildProcess::new(Command::new("cargo").configure( |cmd| { cmd.arg("run") @@ -44,7 +52,7 @@ async fn main() -> Result<()> { ))?) .await?; - // 1) Synchronous call. `quick_echo` has the default task_support = forbidden. + // 1) Synchronous call. let echo = client .call_tool( CallToolRequestParams::new("quick_echo") @@ -53,68 +61,63 @@ async fn main() -> Result<()> { .await?; tracing::info!("quick_echo -> {echo:#?}"); - // 2) Task call. `slow_sum` is task_support = required, so we MUST attach - // `task` metadata. An empty `TaskMetadata` is fine; use `.with_ttl(...)` - // to set a retention window. - let create = client - .send_request(ClientRequest::CallToolRequest(Request::new( - CallToolRequestParams::new("slow_sum") - .with_arguments(object!({ "a": 40, "b": 2 })) - .with_task(TaskMetadata::new()), - ))) + // 2) Task-eligible call. The server sees our tasks capability and + // materializes a task instead of blocking. + let response = client + .call_tool_once( + CallToolRequestParams::new("slow_sum").with_arguments(object!({ "a": 40, "b": 2 })), + ) .await?; - let ServerResult::CreateTaskResult(create) = create else { - return Err(anyhow!("expected CreateTaskResult, got {create:?}")); + let create = match response { + CallToolResponse::Task(create) => create, + CallToolResponse::Complete(result) => { + // The server is allowed to answer synchronously. + tracing::info!("slow_sum answered synchronously -> {result:#?}"); + client.cancel().await?; + return Ok(()); + } + other => return Err(anyhow!("unexpected response: {other:?}")), }; let task_id = create.task.task_id.clone(); + let poll_ms = create.task.poll_interval_ms.unwrap_or(500); tracing::info!( - "slow_sum enqueued as task {task_id} (status = {:?})", + "slow_sum materialized as task {task_id} (status = {:?})", create.task.status ); - // 3) Poll `tasks/get` until the server reports a terminal status. - let final_status = loop { - tokio::time::sleep(std::time::Duration::from_millis(250)).await; + // 3) Poll `tasks/get` until the task reaches a terminal status. + let final_task = loop { + tokio::time::sleep(std::time::Duration::from_millis(poll_ms)).await; let info = client - .send_request(ClientRequest::GetTaskRequest(Request::new( - GetTaskParams::new(task_id.clone()), - ))) + .peer() + .get_task(GetTaskParams::new(task_id.clone())) .await?; - let ServerResult::GetTaskResult(info) = info else { - return Err(anyhow!("expected GetTaskResult, got {info:?}")); - }; - tracing::info!("status = {:?}", info.task.status); + tracing::info!("status = {:?}", info.task.status()); - match info.task.status { - TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Cancelled => { - break info.task.status; - } - _ => {} + if info.task.status().is_terminal() { + break info.task; } }; - if final_status != TaskStatus::Completed { - return Err(anyhow!("task ended in {final_status:?}")); + // The completed task carries the final CallToolResult inline. + match &final_task.payload { + TaskPayload::Completed { result } => { + let call_result: CallToolResult = + serde_json::from_value(serde_json::Value::Object(result.clone()))?; + tracing::info!("slow_sum result -> {call_result:#?}"); + } + TaskPayload::Failed { error } => { + return Err(anyhow!("task failed: {error:?}")); + } + other => { + return Err(anyhow!( + "task ended in unexpected state {:?}", + other.status() + )); + } } - - // 4) Fetch the payload. The server-side handler returns a serialized - // `CallToolResult`. On the wire the response is just a JSON value, and - // `ServerResult` is `#[serde(untagged)]`, so the client decodes it as - // whichever variant the JSON shape matches first — a `CallToolResult` - // here. (For a non-tool task the same value would surface as - // `ServerResult::CustomResult` and need manual `serde_json::from_value`.) - let payload = client - .send_request(ClientRequest::GetTaskPayloadRequest(Request::new( - GetTaskPayloadParams::new(task_id.clone()), - ))) - .await?; - let call_result: CallToolResult = match payload { - ServerResult::CallToolResult(r) => r, - ServerResult::CustomResult(c) => serde_json::from_value(c.0)?, - other => return Err(anyhow!("unexpected task result: {other:?}")), - }; - tracing::info!("slow_sum result -> {call_result:#?}"); + debug_assert_eq!(final_task.status(), TaskStatus::Completed); client.cancel().await?; Ok(()) diff --git a/examples/servers/Cargo.toml b/examples/servers/Cargo.toml index a544f3e..bc983a1 100644 --- a/examples/servers/Cargo.toml +++ b/examples/servers/Cargo.toml @@ -13,6 +13,7 @@ rmcp = { workspace = true, features = [ "transport-streamable-http-server", "auth", "elicitation", + "request-state", "schemars", ] } tokio = { version = "1", features = [ @@ -113,3 +114,11 @@ path = "src/elicitation_enum_inference.rs" [[example]] name = "servers_task_stdio" path = "src/task_stdio.rs" + +[[example]] +name = "servers_mrtr" +path = "src/mrtr.rs" + +[[example]] +name = "servers_subscriptions_streamhttp" +path = "src/subscriptions_streamhttp.rs" diff --git a/examples/servers/README.md b/examples/servers/README.md index 6912651..86bf3a3 100644 --- a/examples/servers/README.md +++ b/examples/servers/README.md @@ -27,6 +27,13 @@ A server using streamable HTTP transport for MCP communication, with axum. - Provides counter tools via HTTP streaming - Demonstrates streamable HTTP transport configuration +### Modern Subscription Server (`subscriptions_streamhttp.rs`) + +A stateless `2026-07-28` server that opens `subscriptions/listen` response +streams, acknowledges the accepted filter, and emits tagged tool-list changes. + +- Run with `cargo run -p mcp-server-examples --example servers_subscriptions_streamhttp` + ### Counter Streamable HTTP Server with Hyper (`counter_hyper_streamable_http.rs`) A server using streamable HTTP transport for MCP communication, with hyper. @@ -64,13 +71,24 @@ A server demonstrating the prompt framework capabilities. ### Task Demo Server (`task_stdio.rs`) -A minimal stdio server demonstrating task-based tool invocation per -[SEP-1319](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks). +A minimal stdio server demonstrating the MCP Tasks extension +([SEP-2663](https://modelcontextprotocol.io/extensions/tasks/overview), `io.modelcontextprotocol/tasks`). -- `slow_sum` is declared with `execution(task_support = "required")`, so clients MUST invoke it as a task +- `slow_sum` is materialized as a task (`CreateTaskResult`, `resultType: "task"`) whenever the client declares the tasks extension capability; other clients get a normal synchronous response - `quick_echo` is a regular synchronous tool for contrast -- Wires up `enqueue_task` / `tasks/get` / `tasks/result` / `tasks/cancel` via `#[task_handler]` -- Pair with `examples/clients/src/task_stdio.rs` to see the full lifecycle (create → poll → fetch result) +- Serves `tasks/get` / `tasks/update` / `tasks/cancel` via a `TaskManager` +- Pair with `examples/clients/src/task_stdio.rs` to see the full lifecycle (create → poll → inline result) + +### MRTR Demo (`mrtr.rs`) + +An end-to-end walkthrough of SEP-2322 Multi Round-Trip Requests, running a +server and client in one process over an in-memory stream. + +- Server answers `tools/call` with an `InputRequiredResult` asking the client to elicit a value +- Client uses `call_tool` to auto-fulfil the elicitation and retry, then `call_tool_once` for manual control +- Seals/opens the untrusted `requestState` with `RequestStateCodec` (HMAC integrity) +- Both sides negotiate `2026-07-28`, the minimum version for MRTR +- Run with `cargo run -p mcp-server-examples --example servers_mrtr` ### Progress Demo Server (`progress_demo.rs`) diff --git a/examples/servers/src/cimd_auth_streamhttp.rs b/examples/servers/src/cimd_auth_streamhttp.rs index 6a634d8..a812977 100644 --- a/examples/servers/src/cimd_auth_streamhttp.rs +++ b/examples/servers/src/cimd_auth_streamhttp.rs @@ -151,16 +151,16 @@ async fn fetch_and_validate_client_metadata(client_id_url: &str) -> Result, -} - -impl OperationResultTransport for ToolCallOperationResult { - fn operation_id(&self) -> &String { - &self.id - } - - fn as_any(&self) -> &dyn Any { - self - } -} - #[derive(Debug, serde::Deserialize, schemars::JsonSchema)] pub struct StructRequest { pub a: i32, @@ -78,7 +61,6 @@ pub struct Counter { counter: Arc>, tool_router: ToolRouter, prompt_router: PromptRouter, - processor: Arc>, } #[tool_router] @@ -89,7 +71,6 @@ impl Counter { counter: Arc::new(Mutex::new(0)), tool_router: Self::tool_router(), prompt_router: Self::prompt_router(), - processor: Arc::new(Mutex::new(OperationProcessor::new())), } } @@ -123,10 +104,7 @@ impl Counter { )])) } - #[tool( - description = "Long running task example", - execution(task_support = "optional") - )] + #[tool(description = "Long running task example")] async fn long_task(&self) -> Result { tokio::time::sleep(std::time::Duration::from_secs(10)).await; Ok(CallToolResult::success(vec![ContentBlock::text( @@ -179,7 +157,7 @@ impl Counter { /// This is an example prompt that takes one required argument, message #[prompt( name = "example_prompt", - meta = Meta(rmcp::object!({"meta_key": "meta_value"})) + meta = MetaObject(rmcp::object!({"meta_key": "meta_value"})) )] async fn example_prompt( &self, @@ -225,12 +203,11 @@ impl Counter { } } -#[tool_handler(meta = Meta(rmcp::object!({"tool_meta_key": "tool_meta_value"})))] -#[prompt_handler(meta = Meta(rmcp::object!({"router_meta_key": "router_meta_value"})))] -#[task_handler] +#[tool_handler(meta = MetaObject(rmcp::object!({"tool_meta_key": "tool_meta_value"})))] +#[prompt_handler(meta = MetaObject(rmcp::object!({"router_meta_key": "router_meta_value"})))] impl ServerHandler for Counter { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_prompts() .enable_resources() @@ -252,8 +229,7 @@ impl ServerHandler for Counter { self._create_resource_text("str:////Users/to/some/path/", "cwd"), self._create_resource_text("memo://insights", "memo-name"), ], - next_cursor: None, - meta: None, + ..Default::default() }) } @@ -261,22 +237,16 @@ impl ServerHandler for Counter { &self, request: ReadResourceRequestParams, _: RequestContext, - ) -> Result { + ) -> Result { let uri = &request.uri; match uri.as_str() { "str:////Users/to/some/path/" => { let cwd = "/Users/to/some/path/"; - Ok(ReadResourceResult::new(vec![ResourceContents::text( - cwd, - uri.clone(), - )])) + Ok(ReadResourceResult::new(vec![ResourceContents::text(cwd, uri.clone())]).into()) } "memo://insights" => { let memo = "Business Intelligence Memo\n\nAnalysis has revealed 5 key insights ..."; - Ok(ReadResourceResult::new(vec![ResourceContents::text( - memo, - uri.clone(), - )])) + Ok(ReadResourceResult::new(vec![ResourceContents::text(memo, uri.clone())]).into()) } _ => Err(McpError::resource_not_found( "resource_not_found", @@ -293,15 +263,14 @@ impl ServerHandler for Counter { _: RequestContext, ) -> Result { Ok(ListResourceTemplatesResult { - next_cursor: None, resource_templates: Vec::new(), - meta: None, + ..Default::default() }) } async fn initialize( &self, - _request: InitializeRequestParams, + request: InitializeRequestParams, context: RequestContext, ) -> Result { if let Some(http_request_part) = context.extensions.get::() { @@ -309,7 +278,8 @@ impl ServerHandler for Counter { let initialize_uri = &http_request_part.uri; tracing::info!(?initialize_headers, %initialize_uri, "initialize from http server"); } - Ok(self.get_info()) + context.peer.set_peer_info(request.clone()); + self.negotiate_initialize(&request) } } @@ -388,49 +358,4 @@ mod tests { let prompts = router.list_all(); assert_eq!(prompts.len(), 2); } - - #[tokio::test] - async fn test_client_enqueues_long_task() -> anyhow::Result<()> { - let counter = Counter::new(); - let processor = counter.processor.clone(); - let client = TestClient::default(); - - let (server_transport, client_transport) = tokio::io::duplex(4096); - let server_handle = tokio::spawn(async move { - let service = counter.serve(server_transport).await?; - service.waiting().await?; - anyhow::Ok(()) - }); - - let client_service = client.serve(client_transport).await?; - let params = CallToolRequestParams::new("long_task").with_task(TaskMetadata::new()); - let response = client_service - .send_request(ClientRequest::CallToolRequest(Request::new(params.clone()))) - .await?; - - let ServerResult::CreateTaskResult(info) = response else { - panic!("expected task creation result, got {response:?}"); - }; - let task = info.task; - - assert_eq!(task.status, TaskStatus::Working); - // task list should show the task - let tasks = client_service - .send_request(ClientRequest::ListTasksRequest( - RequestOptionalParam::default(), - )) - .await - .unwrap(); - let ServerResult::ListTasksResult(listed) = tasks else { - panic!("expected list tasks result, got {tasks:?}"); - }; - assert_eq!(listed.tasks[0].task_id, task.task_id); - tokio::time::sleep(Duration::from_millis(50)).await; - let running = processor.lock().await.running_task_count(); - assert_eq!(running, 1); - - client_service.cancel().await?; - let _ = server_handle.await; - Ok(()) - } } diff --git a/examples/servers/src/common/progress_demo.rs b/examples/servers/src/common/progress_demo.rs index 253dbc2..68a32ec 100644 --- a/examples/servers/src/common/progress_demo.rs +++ b/examples/servers/src/common/progress_demo.rs @@ -75,13 +75,13 @@ impl ProgressDemo { ctx.meta.get_key_value("progressToken") ); let Some((_, progress_token)) = ctx.meta.get_key_value("progressToken") else { - return Err(McpError::internal_error(format!("No progress token"), None)); + return Err(McpError::internal_error("No progress token", None)); }; let Ok(progress_token) = serde_json::from_value::(progress_token.clone()) else { return Err(McpError::internal_error( - format!("Invalid format of the progress token"), + "Invalid format of the progress token", None, )); }; @@ -131,8 +131,8 @@ impl ProgressDemo { #[tool_handler] impl ServerHandler for ProgressDemo { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) .with_protocol_version(ProtocolVersion::V_2024_11_05) .with_server_info(Implementation::from_build_env()) .with_instructions( diff --git a/examples/servers/src/common/task_demo.rs b/examples/servers/src/common/task_demo.rs index 275047a..30aed17 100644 --- a/examples/servers/src/common/task_demo.rs +++ b/examples/servers/src/common/task_demo.rs @@ -1,27 +1,26 @@ -//! Minimal example of a tool that supports task-based invocation (SEP-1319). +//! Minimal example of a server that supports the MCP Tasks extension +//! (SEP-2663, `io.modelcontextprotocol/tasks`). //! -//! - `slow_sum` is marked `task_support = "required"`, so the client MUST invoke -//! it as a task. The server enqueues the call into an `OperationProcessor`, -//! returns a task id immediately, and the client polls `tasks/get` and -//! fetches the payload via `tasks/result`. -//! - `quick_echo` is a regular synchronous tool for contrast (the default, -//! `task_support = "forbidden"`). +//! - `slow_sum` is executed as a task whenever the client declares the tasks +//! extension capability: the server returns a `CreateTaskResult` +//! (`resultType: "task"`) immediately and the client polls `tasks/get`. +//! Clients that do not declare the extension get a normal synchronous +//! response. +//! - `quick_echo` is a regular synchronous tool for contrast. //! //! See `examples/clients/src/task_stdio.rs` for the matching client. #![allow(dead_code)] -use std::sync::Arc; - use rmcp::{ ErrorData as McpError, ServerHandler, handler::server::{router::tool::ToolRouter, wrapper::Parameters}, - model::{CallToolResult, ContentBlock}, - schemars, task_handler, - task_manager::OperationProcessor, - tool, tool_handler, tool_router, + model::*, + schemars, + service::{RequestContext, RoleServer}, + task_manager::{TaskExit, TaskManager, TaskOptions}, + tool, tool_router, }; -use tokio::sync::Mutex; #[derive(Debug, serde::Deserialize, schemars::JsonSchema)] pub struct SumArgs { @@ -34,13 +33,10 @@ pub struct EchoArgs { pub message: String, } -/// Server state. The `processor` field is required by `#[task_handler]`: -/// the macro generates `enqueue_task` / `tasks/*` handlers that submit and -/// poll operations through it. #[derive(Clone)] pub struct TaskDemo { tool_router: ToolRouter, - processor: Arc>, + tasks: TaskManager, } impl Default for TaskDemo { @@ -54,17 +50,12 @@ impl TaskDemo { pub fn new() -> Self { Self { tool_router: Self::tool_router(), - processor: Arc::new(Mutex::new(OperationProcessor::new())), + tasks: TaskManager::new(), } } - /// Long-running tool. The `execution(task_support = "required")` attribute - /// tells clients they MUST call this tool as a task; the server returns - /// `-32601` if they don't. - #[tool( - description = "Sum two numbers after a 2-second delay", - execution(task_support = "required") - )] + /// Long-running tool. Run as a task when the client supports tasks. + #[tool(description = "Sum two numbers after a 2-second delay")] async fn slow_sum( &self, Parameters(SumArgs { a, b }): Parameters, @@ -75,7 +66,7 @@ impl TaskDemo { )])) } - /// Synchronous tool with the default `task_support = "forbidden"`. + /// Synchronous tool. #[tool(description = "Echo a message back immediately")] async fn quick_echo( &self, @@ -85,9 +76,92 @@ impl TaskDemo { } } -/// `#[task_handler]` reads `self.processor` (configurable via the macro's -/// `processor = ...` argument) and synthesizes `enqueue_task`, `list_tasks`, -/// `get_task_info`, `get_task_result`, and `cancel_task` for us. -#[tool_handler] -#[task_handler] -impl ServerHandler for TaskDemo {} +impl ServerHandler for TaskDemo { + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + // SEP-2663: the server decides per-request whether to materialize a + // task, but MUST NOT return one unless the request declared the tasks + // extension capability. + let client_supports_tasks = context + .client_capabilities() + .is_some_and(|caps| caps.supports_tasks()); + + if request.name == "slow_sum" && client_supports_tasks { + let params: SumArgs = serde_json::from_value(serde_json::Value::Object( + request.arguments.clone().unwrap_or_default(), + )) + .map_err(|e| McpError::invalid_params(e.to_string(), None))?; + let task = self.tasks.spawn(TaskOptions::default(), move |ctx| { + Box::pin(async move { + // Cancellation is cooperative (SEP-2663): honor + // tasks/cancel by exiting with TaskExit::Cancelled. + tokio::select! { + _ = ctx.cancelled() => { + Err(TaskExit::Cancelled) + } + _ = tokio::time::sleep(std::time::Duration::from_secs(2)) => { + Ok(CallToolResult::success(vec![ContentBlock::text( + (params.a + params.b).to_string(), + )])) + } + } + }) + }); + return Ok(CallToolResponse::Task(CreateTaskResult::new(task))); + } + + // Fall back to synchronous execution via the tool router. + let tcc = rmcp::handler::server::tool::ToolCallContext::new(self, request, context); + self.tool_router.call(tcc).await + } + + async fn list_tools( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + Ok(ListToolsResult::with_all_items(self.tool_router.list_all())) + } + + async fn get_task( + &self, + request: GetTaskParams, + _context: RequestContext, + ) -> Result { + Ok(GetTaskResult::new(self.tasks.get_task(&request.task_id)?)) + } + + async fn update_task( + &self, + request: UpdateTaskParams, + _context: RequestContext, + ) -> Result<(), McpError> { + self.tasks + .update_task(&request.task_id, request.input_responses) + } + + async fn cancel_task( + &self, + request: CancelTaskParams, + _context: RequestContext, + ) -> Result<(), McpError> { + self.tasks.cancel_task(&request.task_id) + } + + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tasks() + .build(), + ) + .with_instructions( + "Task demo server (SEP-2663). `slow_sum` runs as a task for \ + clients that declare the tasks extension." + .to_string(), + ) + } +} diff --git a/examples/servers/src/completion_stdio.rs b/examples/servers/src/completion_stdio.rs index 812ed31..c7894ea 100644 --- a/examples/servers/src/completion_stdio.rs +++ b/examples/servers/src/completion_stdio.rs @@ -123,23 +123,23 @@ impl SqlQueryServer { .collect(); // If no uppercase letters found, just use first letter - if first_chars.is_empty() && !candidate.is_empty() { - if let Some(first) = candidate.chars().next() { - first_chars.push(first.to_lowercase().next().unwrap_or('\0')); - } + if first_chars.is_empty() + && let Some(first) = candidate.chars().next() + { + first_chars.push(first.to_lowercase().next().unwrap_or('\0')); } } // Special case: if query is 2 chars and we only got 1 char, try matching first 2 letters - if query_chars.len() == 2 && first_chars.len() == 1 { - if let Some(first) = candidate.chars().nth(0) { - if let Some(second) = candidate.chars().nth(1) { - first_chars = vec![ - first.to_lowercase().next().unwrap_or('\0'), - second.to_lowercase().next().unwrap_or('\0'), - ]; - } - } + if query_chars.len() == 2 + && first_chars.len() == 1 + && let Some(first) = candidate.chars().next() + && let Some(second) = candidate.chars().nth(1) + { + first_chars = vec![ + first.to_lowercase().next().unwrap_or('\0'), + second.to_lowercase().next().unwrap_or('\0'), + ]; } if query_chars.len() != first_chars.len() { @@ -193,7 +193,7 @@ impl SqlQueryServer { } } -#[prompt_router] +#[prompt_router(router = "prompt_router")] impl SqlQueryServer { #[prompt(name = "sql_query", description = "Smart SQL query builder")] async fn sql_query( @@ -308,10 +308,10 @@ impl SqlQueryServer { } } -#[prompt_handler] +#[prompt_handler(router = self.prompt_router)] impl ServerHandler for SqlQueryServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_completions() .enable_prompts() diff --git a/examples/servers/src/complex_auth_streamhttp.rs b/examples/servers/src/complex_auth_streamhttp.rs index 34c4b15..bf763d3 100644 --- a/examples/servers/src/complex_auth_streamhttp.rs +++ b/examples/servers/src/complex_auth_streamhttp.rs @@ -69,10 +69,10 @@ impl McpOAuthStore { redirect_uri: &str, ) -> Option { let clients = self.clients.read().await; - if let Some(client) = clients.get(client_id) { - if client.redirect_uri.contains(&redirect_uri.to_string()) { - return Some(client.clone()); - } + if let Some(client) = clients.get(client_id) + && client.redirect_uri == redirect_uri + { + return Some(client.clone()); } None } diff --git a/examples/servers/src/elicitation_enum_inference.rs b/examples/servers/src/elicitation_enum_inference.rs index bed5c1d..e565892 100644 --- a/examples/servers/src/elicitation_enum_inference.rs +++ b/examples/servers/src/elicitation_enum_inference.rs @@ -97,7 +97,7 @@ struct ElicitationEnumFormServer { tool_router: ToolRouter, } -#[tool_router] +#[tool_router(router = tool_router)] impl ElicitationEnumFormServer { pub fn new() -> Self { Self { @@ -153,10 +153,10 @@ impl ElicitationEnumFormServer { } } -#[tool_handler] +#[tool_handler(router = self.tool_router)] impl ServerHandler for ElicitationEnumFormServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) .with_server_info(Implementation::from_build_env()) .with_instructions( "Simple server demonstrating elicitation for enum selection".to_string(), diff --git a/examples/servers/src/elicitation_stdio.rs b/examples/servers/src/elicitation_stdio.rs index e465e4b..5d7f280 100644 --- a/examples/servers/src/elicitation_stdio.rs +++ b/examples/servers/src/elicitation_stdio.rs @@ -64,7 +64,7 @@ impl Default for ElicitationServer { } } -#[tool_router] +#[tool_router(router = tool_router)] impl ElicitationServer { #[tool(description = "Greet user with name collection")] async fn greet_user( @@ -129,18 +129,9 @@ impl ElicitationServer { ) })?; match elicit_result { - ElicitationAction::Accept => { - // Mock notifying completion - let _ = context - .peer - .notify_url_elicitation_completed(ElicitationResponseNotificationParam::new( - "elicit_123", - )) - .await; - Ok(CallToolResult::success(vec![ContentBlock::text( - "Elicitation via URL successful".to_string(), - )])) - } + ElicitationAction::Accept => Ok(CallToolResult::success(vec![ContentBlock::text( + "Elicitation via URL successful".to_string(), + )])), ElicitationAction::Cancel => Ok(CallToolResult::success(vec![ContentBlock::text( "Elicitation via URL cancelled by user".to_string(), )])), @@ -154,10 +145,10 @@ impl ElicitationServer { } } -#[tool_handler] +#[tool_handler(router = self.tool_router)] impl ServerHandler for ElicitationServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) .with_server_info(Implementation::from_build_env()) .with_instructions( "Simple server demonstrating elicitation for user name collection".to_string(), diff --git a/examples/servers/src/mrtr.rs b/examples/servers/src/mrtr.rs new file mode 100644 index 0000000..b6c87b6 --- /dev/null +++ b/examples/servers/src/mrtr.rs @@ -0,0 +1,333 @@ +//! SEP-2322 Multi Round-Trip Request (MRTR) end-to-end example. +//! +//! This runs a server and a client in the same process, connected over an +//! in-memory duplex stream, to show the full MRTR flow: +//! +//! * The **server** answers `tools/call` with an [`InputRequiredResult`] instead +//! of a final result, asking the client to elicit a value first. It stores its +//! progress in an opaque, integrity-protected `requestState` produced by a +//! [`RequestStateCodec`], and verifies that state when the client retries. +//! * The **client** uses the high-level [`RunningService::call_tool`] helper, +//! which automatically fulfils the elicitation through the local +//! [`ClientHandler`] and retries the original request. The example then repeats +//! the call with [`RunningService::call_tool_once`] to show the manual escape +//! hatch that returns the intermediate result without retrying. +//! +//! ## Version gating +//! +//! `InputRequiredResult` is only valid once the peers have negotiated protocol +//! version `2026-07-28` or newer. Both sides advertise that version below. If a +//! server emits an `InputRequiredResult` to an older client, the SDK turns it +//! into a protocol error instead of sending it on the wire. +//! +//! ## `requestState` is untrusted input +//! +//! The client echoes `requestState` back verbatim, so from the server's point of +//! view it is attacker-controlled. A stateless server that puts meaningful data +//! in `requestState` MUST verify it. This example uses [`RequestStateCodec`] to +//! seal and open it with an HMAC tag; tampered values are rejected. +//! +//! Run with: +//! +//! ```sh +//! cargo run -p mcp-server-examples --example servers_mrtr +//! ``` + +use rmcp::{ + ClientHandler, ServerHandler, ServiceExt, + model::*, + service::{RequestContext, RoleClient, RoleServer}, +}; +use serde::{Deserialize, Serialize}; +use serde_json::json; + +/// A fixed key for this example. Real deployments should load at least 32 bytes +/// of random secret key material from a secret manager and keep it stable for +/// the lifetime of any in-flight MRTR exchange. +const REQUEST_STATE_KEY: &[u8] = b"example-request-state-signing-key-32b!"; + +#[derive(Debug, Serialize, Deserialize)] +struct WeatherRequestState { + awaiting: WeatherInput, +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +enum WeatherInput { + City, +} + +/// A server that needs a city name before it can answer a weather query. +#[derive(Clone)] +struct WeatherServer { + codec: RequestStateCodec, +} + +impl WeatherServer { + fn new(key: impl Into>) -> Result { + Ok(Self { + codec: RequestStateCodec::try_new(key)?, + }) + } +} + +impl WeatherServer { + fn open_request_state(&self, sealed: &str) -> Result { + self.codec + .open_json(sealed) + .map_err(|_| ErrorData::invalid_params("tampered or unknown request state", None)) + } +} + +fn finish_weather_request( + input_responses: Option<&InputResponses>, +) -> Result { + let response = input_responses + .and_then(|responses| responses.get("city")) + .ok_or_else(|| ErrorData::invalid_params("missing city elicitation response", None))?; + let response = ElicitResult::deserialize(response) + .map_err(|_| ErrorData::invalid_params("invalid city elicitation response", None))?; + + let message = match response.action { + ElicitationAction::Accept => { + let city = response + .content + .as_ref() + .and_then(|content| content["city"].as_str()) + .ok_or_else(|| { + ErrorData::invalid_params("accepted elicitation did not include a city", None) + })?; + format!("It is sunny in {city}.") + } + ElicitationAction::Decline => "I cannot look up the weather without a city.".into(), + ElicitationAction::Cancel => "Weather lookup cancelled.".into(), + _ => { + return Err(ErrorData::invalid_params( + "unsupported elicitation action", + None, + )); + } + }; + + Ok(CallToolResult::success(vec![ContentBlock::text(message)])) +} + +impl ServerHandler for WeatherServer { + fn get_info(&self) -> ServerConfig { + let mut info = ServerConfig::new(ServerCapabilities::builder().enable_tools().build()); + // MRTR requires 2026-07-28 or newer. + info.protocol_version = ProtocolVersion::V_2026_07_28; + info + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + match request.request_state { + // First round: ask the client to provide a city, and remember where + // we are by sealing our progress into `requestState`. + None => { + let sealed = self + .codec + .seal_json(&WeatherRequestState { + awaiting: WeatherInput::City, + }) + .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; + + let mut input_requests = InputRequests::new(); + input_requests.insert( + "city".to_string(), + InputRequest::Elicitation(ElicitRequest::new( + ElicitRequestParams::FormElicitationParams { + meta: None, + message: "Which city do you want the weather for?".into(), + requested_schema: serde_json::from_value(json!({ + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"] + })) + .expect("valid schema"), + }, + )), + ); + + Ok(InputRequiredResult::new(Some(input_requests), Some(sealed)).into()) + } + // Retry round: verify the echoed state before trusting it, read the + // elicited city, and return the final result. + Some(sealed) => { + let state = self.open_request_state(&sealed)?; + match state.awaiting { + WeatherInput::City => finish_weather_request(request.input_responses.as_ref()) + .map(CallToolResponse::from), + } + } + } + } +} + +/// A client that fulfils elicitation requests. A real client would prompt a user. +#[derive(Clone, Default)] +struct InteractiveClient; + +impl ClientHandler for InteractiveClient { + fn get_info(&self) -> ClientConfig { + ClientConfig::new( + ClientCapabilities::builder().enable_elicitation().build(), + Implementation::new("mrtr-example-client", env!("CARGO_PKG_VERSION")), + ) + .with_protocol_version(ProtocolVersion::V_2026_07_28) + } + + async fn create_elicitation( + &self, + request: ElicitRequestParams, + _context: RequestContext, + ) -> Result { + if let ElicitRequestParams::FormElicitationParams { message, .. } = &request { + println!(" [client] server asked: {message}"); + } + // Pretend the user typed "Paris". + Ok(ElicitResult::new(ElicitationAction::Accept).with_content(json!({ "city": "Paris" }))) + } +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let (server_transport, client_transport) = tokio::io::duplex(8192); + let server = WeatherServer::new(REQUEST_STATE_KEY)?; + + // Spin up the server side. + tokio::spawn(async move { + let server = server + .serve(server_transport) + .await + .expect("server should start"); + let _ = server.waiting().await; + }); + + // Connect the client (this performs the initialize handshake). + let client = InteractiveClient.serve(client_transport).await?; + + // 1. High-level auto mode: the SDK fulfils the elicitation and retries for us. + println!("== auto mode (call_tool) =="); + let result = client + .call_tool(CallToolRequestParams::new("weather")) + .await?; + println!( + " [client] final result: {}\n", + result.content[0].as_text().unwrap().text + ); + + // 2. Manual mode: get the intermediate InputRequiredResult without retrying. + println!("== manual mode (call_tool_once) =="); + match client + .call_tool_once(CallToolRequestParams::new("weather")) + .await? + { + CallToolResponse::InputRequired(input_required) => { + let requests = input_required.input_requests.unwrap_or_default(); + println!( + " [client] server requested {} input(s); handling them yourself is up to you.", + requests.len() + ); + } + CallToolResponse::Complete(result) => { + println!(" [client] completed immediately: {result:?}"); + } + _ => println!(" [client] unhandled response variant"), + } + + client.cancel().await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn input_responses(response: ElicitResult) -> InputResponses { + [("city".into(), serde_json::to_value(response).unwrap())].into() + } + + #[test] + fn finish_weather_request_uses_city_when_elicitation_is_accepted() { + let responses = input_responses( + ElicitResult::new(ElicitationAction::Accept).with_content(json!({ "city": "Paris" })), + ); + + let result = finish_weather_request(Some(&responses)).unwrap(); + + assert_eq!( + result.content[0].as_text().unwrap().text, + "It is sunny in Paris." + ); + } + + #[test] + fn finish_weather_request_ignores_content_when_elicitation_is_declined() { + let responses = input_responses( + ElicitResult::new(ElicitationAction::Decline).with_content(json!({ "city": "Paris" })), + ); + + let result = finish_weather_request(Some(&responses)).unwrap(); + + assert_eq!( + result.content[0].as_text().unwrap().text, + "I cannot look up the weather without a city." + ); + } + + #[test] + fn finish_weather_request_stops_lookup_when_elicitation_is_cancelled() { + let responses = input_responses(ElicitResult::new(ElicitationAction::Cancel)); + + let result = finish_weather_request(Some(&responses)).unwrap(); + + assert_eq!( + result.content[0].as_text().unwrap().text, + "Weather lookup cancelled." + ); + } + + #[test] + fn finish_weather_request_rejects_accepted_elicitation_without_city() { + let responses = + input_responses(ElicitResult::new(ElicitationAction::Accept).with_content(json!({}))); + + let error = finish_weather_request(Some(&responses)).unwrap_err(); + + assert_eq!( + error, + ErrorData::invalid_params("accepted elicitation did not include a city", None) + ); + } + + #[test] + fn finish_weather_request_rejects_missing_city_response() { + let error = finish_weather_request(None).unwrap_err(); + + assert_eq!( + error, + ErrorData::invalid_params("missing city elicitation response", None) + ); + } + + #[test] + fn open_request_state_rejects_unknown_phase() { + let server = WeatherServer::new(REQUEST_STATE_KEY).unwrap(); + let sealed = server + .codec + .seal_json(&json!({ "awaiting": "country" })) + .unwrap(); + + let error = server.open_request_state(&sealed).unwrap_err(); + + assert_eq!( + error, + ErrorData::invalid_params("tampered or unknown request state", None) + ); + } +} diff --git a/examples/servers/src/prompt_stdio.rs b/examples/servers/src/prompt_stdio.rs index 7b6e285..7f018d0 100644 --- a/examples/servers/src/prompt_stdio.rs +++ b/examples/servers/src/prompt_stdio.rs @@ -112,7 +112,7 @@ impl Default for PromptServer { } } -#[prompt_router] +#[prompt_router(router = "prompt_router")] impl PromptServer { /// Simple greeting prompt without parameters #[prompt( @@ -305,17 +305,17 @@ impl PromptServer { ]; // Add tried solutions if any - if let Some(tried) = args.tried_solutions { - if !tried.is_empty() { - messages.push(PromptMessage::new_text( - Role::User, - format!("I've already tried: {}", tried.join(", ")), - )); - messages.push(PromptMessage::new_text( - Role::Assistant, - "I see you've already attempted some solutions. Let me suggest different approaches.", - )); - } + if let Some(tried) = args.tried_solutions + && !tried.is_empty() + { + messages.push(PromptMessage::new_text( + Role::User, + format!("I've already tried: {}", tried.join(", ")), + )); + messages.push(PromptMessage::new_text( + Role::Assistant, + "I see you've already attempted some solutions. Let me suggest different approaches.", + )); } messages.push(PromptMessage::new_text( @@ -361,10 +361,10 @@ impl PromptServer { } } -#[prompt_handler] +#[prompt_handler(router = self.prompt_router)] impl ServerHandler for PromptServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_prompts().build()).with_instructions( + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_prompts().build()).with_instructions( "This server provides various prompt templates for code review, data analysis, \ writing assistance, debugging help, and personalized learning paths. \ All prompts are designed to provide structured, context-aware assistance.", diff --git a/examples/servers/src/sampling_stdio.rs b/examples/servers/src/sampling_stdio.rs index f75c445..41f74c7 100644 --- a/examples/servers/src/sampling_stdio.rs +++ b/examples/servers/src/sampling_stdio.rs @@ -1,4 +1,7 @@ -#![allow(deprecated)] +#![expect( + deprecated, + reason = "This example demonstrates the deprecated MCP sampling API" +)] use std::sync::Arc; use anyhow::Result; @@ -18,8 +21,8 @@ use tracing_subscriber::{self, EnvFilter}; pub struct SamplingDemoServer; impl ServerHandler for SamplingDemoServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) .with_instructions(concat!( "This is a demo server that requests sampling from clients. It provides tools that use LLM capabilities.\n\n", "IMPORTANT: This server requires a client that supports the 'sampling/createMessage' method. ", @@ -31,7 +34,7 @@ impl ServerHandler for SamplingDemoServer { &self, request: CallToolRequestParams, context: RequestContext, - ) -> Result { + ) -> Result { match request.name.as_ref() { "ask_llm" => { // Get the question from arguments @@ -79,7 +82,8 @@ impl ServerHandler for SamplingDemoServer { .and_then(|c| c.as_text()) .map(|t| &t.text) .unwrap_or(&"No text response".to_string()) - ))])) + ))]) + .into()) } _ => Err(ErrorData::new( @@ -113,8 +117,7 @@ impl ServerHandler for SamplingDemoServer { .unwrap(), ), )], - meta: None, - next_cursor: None, + ..Default::default() }) } } diff --git a/examples/servers/src/subscriptions_streamhttp.rs b/examples/servers/src/subscriptions_streamhttp.rs new file mode 100644 index 0000000..7615d59 --- /dev/null +++ b/examples/servers/src/subscriptions_streamhttp.rs @@ -0,0 +1,83 @@ +use std::{borrow::Cow, time::Duration}; + +use rmcp::{ + ErrorData, ServerHandler, + model::{ProtocolVersion, ServerCapabilities, ServerConfig, SubscriptionFilter}, + service::SubscriptionContext, + transport::streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, + }, +}; +use tokio_util::sync::CancellationToken; + +#[derive(Clone)] +struct SubscriptionServer; + +impl ServerHandler for SubscriptionServer { + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(&[ProtocolVersion::V_2026_07_28]) + } + + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_tool_list_changed() + .build(), + ) + } + + fn accepted_subscription_filter( + &self, + requested: &SubscriptionFilter, + ) -> Option { + Some(requested.supported_by(&self.get_info().capabilities)) + } + + async fn listen(&self, context: SubscriptionContext) -> Result<(), ErrorData> { + loop { + tokio::select! { + () = context.cancelled() => return Ok(()), + () = tokio::time::sleep(Duration::from_secs(2)) => { + context + .sink() + .notify_tool_list_changed() + .await + .map_err(|error| { + ErrorData::internal_error(error.to_string(), None) + })?; + } + } + } + } +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()), + ) + .init(); + + let cancellation_token = CancellationToken::new(); + let service: StreamableHttpService = + StreamableHttpService::new( + || Ok(SubscriptionServer), + Default::default(), + StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_sse_keep_alive(Some(Duration::from_secs(10))) + .with_cancellation_token(cancellation_token.child_token()), + ); + let router = axum::Router::new().nest_service("/mcp", service); + let listener = tokio::net::TcpListener::bind("127.0.0.1:8000").await?; + + axum::serve(listener, router) + .with_graceful_shutdown(async move { + let _ = tokio::signal::ctrl_c().await; + cancellation_token.cancel(); + }) + .await?; + Ok(()) +} diff --git a/examples/transport/Cargo.toml b/examples/transport/Cargo.toml index 716b322..6ffe963 100644 --- a/examples/transport/Cargo.toml +++ b/examples/transport/Cargo.toml @@ -1,6 +1,7 @@ [package] name = "transport" edition = { workspace = true } +rust-version = { workspace = true } version = { workspace = true } authors = { workspace = true } license = { workspace = true } @@ -40,7 +41,7 @@ rand = { version = "0.10" } schemars = { version = "1.0", optional = true } hyper = { version = "1", features = ["client", "server", "http1"] } hyper-util = { version = "0.1", features = ["tokio"] } -tokio-tungstenite = "0.29.0" +tokio-tungstenite = "0.30.0" reqwest = { version = "0.13.2" } pin-project-lite = "0.2" diff --git a/examples/wasi/Cargo.toml b/examples/wasi/Cargo.toml index 4cceb43..4eb7735 100644 --- a/examples/wasi/Cargo.toml +++ b/examples/wasi/Cargo.toml @@ -1,6 +1,7 @@ [package] name = "wasi-mcp-example" edition = { workspace = true } +rust-version = { workspace = true } version = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/scripts/keycloak-oauth-fixture.sh b/scripts/keycloak-oauth-fixture.sh new file mode 100755 index 0000000..50c9b85 --- /dev/null +++ b/scripts/keycloak-oauth-fixture.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# ============================================================================= +# keycloak-oauth-fixture.sh — Live OAuth fixture for rmcp refresh tests +# +# Starts a throwaway Keycloak and provisions the realm that +# crates/rmcp/tests/test_live_oauth_refresh.rs expects. Those tests are +# #[ignore]d, so nothing here runs in CI. +# +# Provision: ./scripts/keycloak-oauth-fixture.sh +# Run tests: cargo test -p rmcp --all-features \ +# --test test_live_oauth_refresh -- --ignored +# Tear down: docker rm -f kc-rmcp-test +# +# Requires Docker. Override the port with KC_BASE (default localhost:8081); +# the tests read the same variable. +# ============================================================================= +set -euo pipefail + +KC=${KC_BASE:-http://localhost:8081} +PORT=${KC##*:} +CONTAINER=kc-rmcp-test + +docker rm -f "$CONTAINER" >/dev/null 2>&1 || true +docker run -d --name "$CONTAINER" -p "$PORT:8080" \ + -e KC_BOOTSTRAP_ADMIN_USERNAME=admin -e KC_BOOTSTRAP_ADMIN_PASSWORD=admin \ + quay.io/keycloak/keycloak:26.0 start-dev >/dev/null + +echo "waiting for keycloak at $KC ..." +until curl -sf -o /dev/null "$KC/realms/master/.well-known/openid-configuration"; do + sleep 3 +done + +token=$(curl -s -X POST "$KC/realms/master/protocol/openid-connect/token" \ + -d client_id=admin-cli -d username=admin -d password=admin -d grant_type=password | + python3 -c 'import sys,json;print(json.load(sys.stdin)["access_token"])') + +provision() { + curl -s -X POST "$KC/admin/realms$1" \ + -H "Authorization: Bearer $token" -H "Content-Type: application/json" \ + -d "$2" -o /dev/null -w " ${1:-/} -> %{http_code}\n" +} + +# revokeRefreshToken makes refresh tokens single-use. The concurrency test +# depends on it: without the refresh guard the second caller replays a consumed +# token and Keycloak answers invalid_grant. +provision "" '{"realm":"rmcp","enabled":true,"revokeRefreshToken":true,"refreshTokenMaxReuse":0}' +provision "/rmcp/clients" '{"clientId":"rmcp-client","secret":"rmcp-secret","publicClient":false, + "directAccessGrantsEnabled":true,"standardFlowEnabled":true, + "redirectUris":["http://localhost/callback"]}' +# The profile fields and empty requiredActions keep the direct access grant from +# failing with "Account is not fully set up". +provision "/rmcp/users" '{"username":"alice","enabled":true,"emailVerified":true, + "email":"alice@example.com","firstName":"Alice","lastName":"Example","requiredActions":[], + "credentials":[{"type":"password","value":"alice-pw","temporary":false}]}' + +echo "ready"