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 9850daf..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@v7 + - 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 9dd9604..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@v7 + 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 @@ -201,9 +210,9 @@ jobs: 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 @@ -211,7 +220,7 @@ jobs: permissions: contents: read steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Read MSRV from Cargo.toml run: | @@ -220,11 +229,11 @@ jobs: echo "MSRV=$MSRV" >> "$GITHUB_ENV" - name: Install MSRV Rust - uses: dtolnay/rust-toolchain@master + uses: dtolnay/rust-toolchain@6c977a6ca4077a0ceb28ffbe03f59d46e9ac8772 # master with: toolchain: ${{ env.MSRV }} - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - name: Check default workspace members with MSRV run: cargo +${{ env.MSRV }} check --all-targets --all-features @@ -233,19 +242,19 @@ jobs: 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@v7 + 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 @@ -253,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 @@ -262,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@v7 + 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 @@ -282,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: | @@ -296,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@v7 + 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 @@ -318,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 @@ -333,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@v7 + 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 @@ -353,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: | @@ -389,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 @@ -406,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: | @@ -432,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 index 1c69e1c..1233deb 100644 --- a/.github/workflows/conformance.yml +++ b/.github/workflows/conformance.yml @@ -12,12 +12,10 @@ concurrency: cancel-in-progress: true env: - # Pinned for reproducible runs; bump deliberately when the suite updates. - CONFORMANCE_VERSION: "0.1.16" - # When updating DRAFT_CONFORMANCE_VERSION, diff - # `conformance list --spec-version 2026-07-28` - # and update #977 - DRAFT_CONFORMANCE_VERSION: "0.2.0-alpha.9" + CONFORMANCE_VERSION: "0.2.0-alpha.10" + +permissions: + contents: read jobs: server: @@ -26,22 +24,20 @@ jobs: permissions: contents: read steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - # Build the whole package (server + client bins): the conformance crate is - # excluded from the workspace default-members. - 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 2025-11-25 server + - name: Start conformance server run: | PORT=8001 ./target/debug/conformance-server & echo $! > server.pid @@ -58,45 +54,19 @@ jobs: 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 + -o conformance-results/2025-11-25 - # These pass today but are excluded from the default "active" suite; - # run them explicitly so regressions are still caught. - - name: Run 2025-11-25 pending scenarios + - name: Run 2026-07-28 server suite run: | - for scenario in json-schema-2020-12 server-sse-polling; do - npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" server \ - --url http://127.0.0.1:8001/mcp \ - --scenario "$scenario" \ - -o conformance-results - done - - - name: Start draft server - run: | - STATELESS=1 PORT=8002 ./target/debug/conformance-server & - echo $! > draft-server.pid - for _ in $(seq 1 30); do - if curl -s -o /dev/null http://127.0.0.1:8002/mcp; then - exit 0 - fi - sleep 1 - done - echo "draft conformance server did not become ready" >&2 - exit 1 - - - name: Run 2026-07-28 versioned-spec server suite - run: | - npx -y "@modelcontextprotocol/conformance@${DRAFT_CONFORMANCE_VERSION}" server \ - --url http://127.0.0.1:8002/mcp \ + 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 + -o conformance-results/2026-07-28 - # Extension scenarios are excluded by the --spec-version filter and - # are informational for tiering. Run them explicitly so their gaps remain - # visible and newly passing scenarios make the strict baseline fail stale. - - name: Run Tasks extension server suite (informational) + - name: Run extension server scenarios (informational) run: | scenarios=( tasks-lifecycle @@ -110,24 +80,21 @@ jobs: tasks-required-task-error tasks-mrtr-composition ) - for scenario in "${scenarios[@]}"; do - npx -y "@modelcontextprotocol/conformance@${DRAFT_CONFORMANCE_VERSION}" server \ - --url http://127.0.0.1:8002/mcp \ + 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 servers + - name: Stop conformance server if: always() - run: | - kill "$(cat draft-server.pid)" 2>/dev/null || true - kill "$(cat server.pid)" 2>/dev/null || true + run: kill "$(cat server.pid)" 2>/dev/null || true - name: Upload results if: always() - uses: actions/upload-artifact@v7 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: conformance-server-results path: | @@ -140,12 +107,12 @@ jobs: permissions: contents: read steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - - uses: Swatinem/rust-cache@v2 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 - name: Build conformance binaries run: cargo build -p mcp-conformance @@ -156,19 +123,19 @@ jobs: --command "$(pwd)/target/debug/conformance-client" \ --suite all \ --spec-version 2025-11-25 \ - -o conformance-client-results/full + -o conformance-client-results/2025-11-25 - - name: Run 2026-07-28 versioned-spec client suite + - name: Run 2026-07-28 client suite run: | - npx -y "@modelcontextprotocol/conformance@${DRAFT_CONFORMANCE_VERSION}" client \ + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" client \ --command "$(pwd)/target/debug/conformance-client" \ --suite all \ --spec-version 2026-07-28 \ - -o conformance-client-results/draft + -o conformance-client-results/2026-07-28 - name: Run extension client suite (informational) run: | - npx -y "@modelcontextprotocol/conformance@${DRAFT_CONFORMANCE_VERSION}" client \ + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" client \ --command "$(pwd)/target/debug/conformance-client" \ --suite extensions \ --expected-failures conformance/expected-failures-extensions.yaml \ @@ -176,7 +143,7 @@ jobs: - name: Upload results if: always() - uses: actions/upload-artifact@v7 + 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 index 3c7d1d7..fe1baa9 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -4,6 +4,9 @@ on: - cron: "0 0 * * *" workflow_dispatch: +permissions: + contents: read + jobs: stale: name: Clean Up @@ -12,7 +15,7 @@ jobs: issues: write pull-requests: write steps: - - uses: actions/stale@v10 + - uses: actions/stale@4391f3da665fdf50b6810c1a66712fb9ba21aa93 # v11 with: days-before-stale: 60 days-before-close: 14 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 821d90c..0629e82 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,13 +4,13 @@ default-members = ["crates/rmcp", "crates/rmcp-macros"] resolver = "2" [workspace.dependencies] -rmcp = { version = "3.0.0-beta.2", path = "./crates/rmcp" } -rmcp-macros = { version = "3.0.0-beta.2", 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" rust-version = "1.88" -version = "3.0.0-beta.2" +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 0863fc6..7b24d6a 100644 --- a/README.md +++ b/README.md @@ -6,20 +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) -This SDK tracks the MCP **`2026-07-28`** draft (the current development spec) -while remaining fully compatible with the stable **`2025-11-25`** release and -earlier versions. New `2026-07-28` features — server discovery & negotiation, +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 alongside the stable feature set. For the full MCP specification, see -[modelcontextprotocol.io](https://modelcontextprotocol.io/specification/draft). +below. For the full MCP specification, see +[modelcontextprotocol.io](https://modelcontextprotocol.io/specification/2026-07-28). ## Table of Contents @@ -28,6 +28,7 @@ below alongside the stable feature set. For the full MCP specification, see - [Resources](#resources) - [Prompts](#prompts) - [Sampling](#sampling) +- [Elicitation](#elicitation) - [Roots](#roots) - [Logging](#logging) - [Completions](#completions) @@ -38,6 +39,10 @@ below alongside the stable feature set. For the full MCP specification, see - [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) @@ -94,10 +99,10 @@ Use [`ClientServiceExt::serve_with_lifecycle`](crates/rmcp/src/service/client.rs select another lifecycle explicitly: ```rust, ignore -use rmcp::{ClientInfo, ClientLifecycleMode, ClientServiceExt, ProtocolVersion}; +use rmcp::{ClientLifecycleMode, ClientServiceExt, ProtocolVersion, model::ClientConfig}; // Start directly with server/discover and include client metadata on every request. -let client = ClientInfo::default() +let client = ClientConfig::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { @@ -107,8 +112,8 @@ let client = ClientInfo::default() .await?; // Or probe the discover lifecycle and fall back when a legacy server reports -// that server/discover is not implemented. -let client = ClientInfo::default() +// that server/discover is not implemented or does not respond within 10 seconds. +let client = ClientConfig::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Auto { @@ -185,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/draft/server/tools) +**MCP Spec:** [Tools](https://modelcontextprotocol.io/specification/2026-07-28/server/tools) ### Server-side @@ -254,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 @@ -274,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/draft/server/resources) +**MCP Spec:** [Resources](https://modelcontextprotocol.io/specification/2026-07-28/server/resources) ### Server-side @@ -293,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(), @@ -328,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 })), @@ -340,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, }) @@ -362,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 @@ -409,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/draft/server/prompts) +**MCP Spec:** [Prompts](https://modelcontextprotocol.io/specification/2026-07-28/server/prompts) ### Server-side @@ -476,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()) } } ``` @@ -487,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 @@ -523,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/draft/client/sampling) +**MCP Spec:** [Sampling](https://modelcontextprotocol.io/specification/2026-07-28/client/sampling) ### Server-side (requesting sampling) @@ -589,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/draft/client/roots) +**MCP Spec:** [Roots](https://modelcontextprotocol.io/specification/2026-07-28/client/roots) ### Server-side @@ -660,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/draft/server/utilities/logging) +**MCP Spec:** [Logging](https://modelcontextprotocol.io/specification/2026-07-28/server/utilities/logging) ### Server-side @@ -670,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(), @@ -733,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/draft/server/utilities/completion) +**MCP Spec:** [Completions](https://modelcontextprotocol.io/specification/2026-07-28/server/utilities/completion) ### Server-side @@ -743,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() @@ -758,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"], @@ -778,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![], }; @@ -799,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) @@ -815,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/draft/basic#notifications) +**MCP Spec:** [Notifications](https://modelcontextprotocol.io/specification/2026-07-28/basic#notifications) ### Progress notifications @@ -865,6 +1151,55 @@ impl ServerHandler for MyServer { } ``` +### Ping + +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 +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. @@ -903,7 +1238,7 @@ Protocol `2026-07-28` replaces `resources/subscribe`, `resources/unsubscribe`, a the standalone HTTP GET stream with the transport-neutral, long-lived `subscriptions/listen` request. Each requested notification category is opt-in. -**MCP Spec:** [Subscriptions](https://modelcontextprotocol.io/specification/draft/basic/patterns/subscriptions) +**MCP Spec:** [Subscriptions](https://modelcontextprotocol.io/specification/2026-07-28/basic/patterns/subscriptions) ### Server-side @@ -918,8 +1253,8 @@ use rmcp::{ }; impl ServerHandler for MyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -995,7 +1330,7 @@ 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/draft/server/tools#multiple-round-trip-requests) +**MCP Spec:** [Multiple Round-Trip Requests](https://modelcontextprotocol.io/specification/2026-07-28/server/tools#multiple-round-trip-requests) ### Server-side @@ -1028,11 +1363,21 @@ async fn call_tool(&self, request: CallToolRequestParams, _ctx: RequestContext **`requestState` is untrusted.** The client echoes it back verbatim, so 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. +> **`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 @@ -1098,7 +1443,7 @@ See [`servers_task_stdio`](examples/servers/src/task_stdio.rs) and the matching ## Caching `rmcp` clients transparently cache responses that carry the -[SEP-2549](https://modelcontextprotocol.io/specification/draft/server/utilities/caching) +[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`. @@ -1150,7 +1495,7 @@ 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/draft/basic/transports#header) +**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`, @@ -1186,7 +1531,7 @@ 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/draft/basic/transports) +**MCP Spec:** [Transports](https://modelcontextprotocol.io/specification/2026-07-28/basic/transports) ### Server-side @@ -1220,6 +1565,13 @@ 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 @@ -1232,13 +1584,231 @@ use rmcp::transport::StreamableHttpClientTransport; // Defaults are stateless-friendly. let transport = StreamableHttpClientTransport::from_uri("http://localhost:8000/mcp"); -let client = ClientInfo::default().serve(transport).await?; +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). @@ -1249,8 +1819,8 @@ See [Oauth_support](docs/OAUTH_SUPPORT.md) for details. ## Related Resources -- [MCP Specification](https://modelcontextprotocol.io/specification/draft) -- [Schema](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/draft/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 60c35b4..6eea344 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -1,116 +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. -Spec 2025-11-25 (suite 0.1.16): Server 100% (30/30) · Client 100% (18/18) -Spec 2026-07-28 (suite 0.2.0-alpha.9): Server 97.5% (39/40) · Client 90.6% (29/32) +**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. -Extension scenarios are reported separately below because they are -informational and do not count toward SDK tiering. +| 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. --- -## Target spec: 2026-07-28 (release 2026-07-28) +## Conformance -All 2026-07-28 work carries the `2026-07-28` label and the -[`2026-07-28 spec` milestone](https://github.com/modelcontextprotocol/rust-sdk/milestone/3). -Per-scenario conformance status is tracked in the epic issue: -[#977 — Tracking: 2026-07-28 spec conformance](https://github.com/modelcontextprotocol/rust-sdk/issues/977). +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). -### Versioned-spec conformance (baseline 2026-07-21, suite `0.2.0-alpha.9`) +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. -- Server: 1 expected failure: `json-schema-2020-12` -- Client: 3 expected failures: `tools_call`, `auth/scope-step-up`, and `auth/authorization-server-migration` -- CI: runs the complete `2026-07-28` versioned-spec suites with a strict baseline; an unlisted failure or a listed scenario that starts passing fails the build +### Informational (not scored for tiering) -### Extension conformance (informational) +Extension-tagged scenarios are excluded by `--spec-version` filters, so CI runs them +in separate steps against `conformance/expected-failures-extensions.yaml`: -Extension-tagged scenarios are excluded by `--spec-version` filters, so CI -runs them in separate server and client steps with -`conformance/expected-failures-extensions.yaml`. - -- SEP-2663 Tasks server: 9 expected failures; `tasks-status-notifications` is currently skipped by the upstream harness; tracked in #868 -- Client extensions: `auth/client-credentials-basic` and `auth/client-credentials-jwt` pass; `auth/enterprise-managed-authorization` is an expected failure +| 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 | ### Spec features without conformance scenarios -Conformance alone does not cover the full spec surface. Feature work tracked via the milestone: +Conformance does not cover the entire spec surface. Remaining feature work tracked via +the milestone: - 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) -(SEP-2575 discovery & negotiation is covered by the `server-stateless` conformance scenario; -implementation is in review — #869, PRs #973, #943.) - -### Release - -The 2026-07-28 implementation ships as **v3.0.0** (release PR #964): MRTR, SEP-2549 cache hints, -SEP-2243 standard headers, and the SEP-2106 relaxations are merged but unreleased — tiering and -relegation are evaluated against the latest stable release, so cutting v3.0.0 with the remaining -conformance fixes is on the critical path. Migration guide (draft, kept current until release): -[discussion #969](https://github.com/modelcontextprotocol/rust-sdk/discussions/969). - ---- - -## Tier 1 (non-conformance requirements) - -### Governance & Policy - -- [ ] Create `VERSIONING.md` — document semver scheme, what constitutes a breaking change, and how breaking changes are communicated -- [ ] Publish a dependency update policy (Tier 1 requires a published policy) -- [ ] Cut v3.0.0 (#964) including all conformance fixes (tier relegation is evaluated against the latest stable release) - -### 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)* - --- ## Completed -- [x] 2025-11-25 server conformance 100% (30 scenarios + pending `json-schema-2020-12`, `server-sse-polling`) -- [x] 2025-11-25 client conformance 100% (18 scenarios + legacy `auth/2025-03-26-*`) -- [x] SEP-2322 MRTR (14 server scenarios + `sep-2322-client-request-state`) +### Tier 1 requirements + +- [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] Cache hints (`caching`) -- [x] `http-header-validation` -- [x] Issue triage labels (bug, enhancement, needs confirmation, needs repro, ready for work, P0–P3) +- [x] SEP-2549 cache hints (`caching`) +- [x] SEP-2243 HTTP standardization (`http-header-validation`, standard headers) +- [x] DNS rebinding protection --- -## Informational (not scored for tiering) +## Nice-to-have (scorecard hygiene, not required for Tier 1) -These extension scenarios are tracked but do not count toward tier advancement: - -| Scenario | Tag | Status | -|---|---|---| -| `auth/client-credentials-jwt` | extension | ✅ Passed | -| `auth/client-credentials-basic` | extension | ✅ Passed | -| `auth/enterprise-managed-authorization` | extension | ❌ Failed — scenario is not implemented by the conformance client | -| `tasks-*` | extension | ❌ 9 expected failures · ⏭️ 1 upstream-skipped scenario | +- [ ] 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/conformance/Cargo.toml b/conformance/Cargo.toml index fc4be6d..e8de6f0 100644 --- a/conformance/Cargo.toml +++ b/conformance/Cargo.toml @@ -19,6 +19,7 @@ rmcp = { path = "../crates/rmcp", features = [ "elicitation", "auth", "auth-client-credentials-jwt", + "auth-enterprise-managed", "request-state", "transport-streamable-http-server", "transport-streamable-http-client-reqwest", @@ -31,6 +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" diff --git a/conformance/expected-failures-extensions.yaml b/conformance/expected-failures-extensions.yaml index fdf9617..0ebbeef 100644 --- a/conformance/expected-failures-extensions.yaml +++ b/conformance/expected-failures-extensions.yaml @@ -1,5 +1,5 @@ # Known failures for informational extension scenarios in -# @modelcontextprotocol/conformance DRAFT_CONFORMANCE_VERSION. +# @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 @@ -9,7 +9,7 @@ # - an unlisted failure fails the build # - a listed scenario that starts passing fails the build as a stale entry # -# When bumping DRAFT_CONFORMANCE_VERSION, review the available extension and +# 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 @@ -21,3 +21,6 @@ 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 852ca09..aef2c27 100644 --- a/conformance/src/bin/client.rs +++ b/conformance/src/bin/client.rs @@ -1,3 +1,10 @@ +#![expect( + deprecated, + reason = "The conformance suite still exercises deprecated sampling scenarios" +)] + +use anyhow::Context; +use oauth2::{ClientSecret, RefreshToken}; use rmcp::{ ClientHandler, ClientLifecycleMode, ClientServiceExt, ErrorData, RoleClient, ServiceExt, model::*, @@ -6,7 +13,8 @@ use rmcp::{ AuthClient, AuthorizationManager, StreamableHttpClientTransport, auth::{ AuthorizationCallback, AuthorizationRequest, ClientCredentialsConfig, - InMemoryCredentialStore, JwtSigningAlgorithm, OAuthState, + InMemoryCredentialStore, JwtSigningAlgorithm, OAuthState, default_oauth_http_client, + enterprise::{EmaAuthorizationServer, EmaClientAuthentication, EmaExchangeRequest}, }, streamable_http_client::StreamableHttpClientTransportConfig, }, @@ -36,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 { @@ -55,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)), @@ -144,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)), @@ -195,25 +214,73 @@ 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( - AuthorizationRequest::new(REDIRECT_URI) - .with_client_name("conformance-client") - .with_client_metadata_url(CIMD_CLIENT_METADATA_URL), - ) + .start_authorization(with_optional_challenge(request, challenge)) .await?; let auth_url = oauth.get_authorization_url().await?; @@ -255,14 +322,14 @@ async fn perform_oauth_flow_preregistered( client_id: &str, client_secret: &str, ) -> anyhow::Result> { + 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_preregistered_client(client_id) + .with_client_secret(client_secret); oauth - .start_authorization( - AuthorizationRequest::new(REDIRECT_URI) - .with_preregistered_client(client_id) - .with_client_secret(client_secret), - ) + .start_authorization(with_optional_challenge(request, challenge)) .await?; let auth_url = oauth.get_authorization_url().await?; @@ -325,14 +392,14 @@ async fn run_auth_scope_step_up_client( server_url: &str, _ctx: &ConformanceContext, ) -> anyhow::Result<()> { + 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( - 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), - ) + .start_authorization(with_optional_challenge(request, challenge)) .await?; let auth_url = oauth.get_authorization_url().await?; @@ -427,15 +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( - AuthorizationRequest::new(REDIRECT_URI) - .with_client_name("conformance-client") - .with_client_metadata_url(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?; @@ -509,7 +576,10 @@ async fn migration_token( return Ok(manager.get_access_token().await?); } - let resolution = manager.resolve_metadata().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, &[]) @@ -609,7 +679,10 @@ async fn run_client_credentials_basic( .unwrap_or("conformance-test-secret"); let mut manager = AuthorizationManager::new(server_url).await?; - let resolution = manager.resolve_metadata().await?; + 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); @@ -706,6 +779,66 @@ async fn run_client_credentials_jwt( Ok(()) } +/// 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?; + + 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?; + } + client.cancel().await?; + Ok(()) +} + /// Cross-app access flow (SEP-1046 extension). async fn run_cross_app_access_client( server_url: &str, @@ -889,8 +1022,7 @@ fn preferred_protocol_versions() -> Vec { preferred_versions } -/// Runs draft stateless scenarios through the public discover lifecycle and -/// Streamable HTTP transport. +/// 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); @@ -1057,6 +1189,11 @@ async fn run_scenario( "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? diff --git a/conformance/src/bin/server.rs b/conformance/src/bin/server.rs index bf902f8..5c4b5c8 100644 --- a/conformance/src/bin/server.rs +++ b/conformance/src/bin/server.rs @@ -52,6 +52,10 @@ fn custom_header_tool() -> Tool { /// 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 { @@ -70,7 +74,7 @@ impl ConformanceServer { 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(REQUEST_STATE_KEY), + request_state_codec: RequestStateCodec::new_unchecked(REQUEST_STATE_KEY), tasks: TaskManager::new(), } } @@ -700,8 +704,8 @@ impl ServerHandler for ConformanceServer { (name == "test_custom_header").then(custom_header_tool) } - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_prompts() .enable_prompts_list_changed() @@ -1740,10 +1744,9 @@ async fn main() -> anyhow::Result<()> { tracing::info!("Starting conformance server on {}", bind_addr); let server = ConformanceServer::new(); - let stateless = std::env::var_os("STATELESS").is_some(); let config = StreamableHttpServerConfig::default() - .with_legacy_session_mode(!stateless) - .with_json_response(stateless); + .with_allowed_origins([format!("http://{bind_addr}")]) + .enforce_origin_validation(); let service = StreamableHttpService::new( move || Ok(server.clone()), LocalSessionManager::default().into(), diff --git a/crates/rmcp-macros/CHANGELOG.md b/crates/rmcp-macros/CHANGELOG.md index 28c40f4..6e35cf3 100644 --- a/crates/rmcp-macros/CHANGELOG.md +++ b/crates/rmcp-macros/CHANGELOG.md @@ -7,6 +7,72 @@ 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 diff --git a/crates/rmcp-macros/Cargo.toml b/crates/rmcp-macros/Cargo.toml index db94afc..9f0b3d3 100644 --- a/crates/rmcp-macros/Cargo.toml +++ b/crates/rmcp-macros/Cargo.toml @@ -16,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/src/lib.rs b/crates/rmcp-macros/src/lib.rs index 156e53b..30433a8 100644 --- a/crates/rmcp-macros/src/lib.rs +++ b/crates/rmcp-macros/src/lib.rs @@ -56,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 /// @@ -122,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()) @@ -185,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()) /// } /// } /// ``` diff --git a/crates/rmcp-macros/src/prompt.rs b/crates/rmcp-macros/src/prompt.rs index 1ebc510..ea35bc1 100644 --- a/crates/rmcp-macros/src/prompt.rs +++ b/crates/rmcp-macros/src/prompt.rs @@ -133,7 +133,7 @@ pub fn prompt(attr: TokenStream, input: TokenStream) -> syn::Result let new_output = syn::parse2::({ let mut lt = quote! { 'static }; if let Some(receiver) = fn_item.sig.receiver() - && let Some((_, receiver_lt)) = receiver.reference.as_ref() + && let syn::ReceiverKind::Reference(_, receiver_lt, _) = &receiver.kind { if let Some(receiver_lt) = receiver_lt { lt = quote! { #receiver_lt }; @@ -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 086eb0d..6d957ec 100644 --- a/crates/rmcp-macros/src/prompt_handler.rs +++ b/crates/rmcp-macros/src/prompt_handler.rs @@ -57,16 +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: None, - cache_scope: 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/tool.rs b/crates/rmcp-macros/src/tool.rs index 69f82e9..c640dd5 100644 --- a/crates/rmcp-macros/src/tool.rs +++ b/crates/rmcp-macros/src/tool.rs @@ -284,7 +284,7 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result { let new_output = syn::parse2::({ let mut lt = quote! { 'static }; if let Some(receiver) = fn_item.sig.receiver() - && let Some((_, receiver_lt)) = receiver.reference.as_ref() + && let syn::ReceiverKind::Reference(_, receiver_lt, _) = &receiver.kind { if let Some(receiver_lt) = receiver_lt { lt = quote! { #receiver_lt }; @@ -342,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 7614668..c9f4a25 100644 --- a/crates/rmcp-macros/src/tool_handler.rs +++ b/crates/rmcp-macros/src/tool_handler.rs @@ -66,15 +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: None, - cache_scope: None, + ttl_ms: supports_cache_hints.then_some(0), + cache_scope: supports_cache_hints + .then_some(rmcp::model::CacheScope::Public), }) } })?; @@ -112,7 +116,7 @@ pub(crate) enum CallerCapability { Prompts, } -/// Build a `get_info()` method that returns `ServerInfo` with the appropriate capabilities. +/// Build a `get_info()` method that returns `ServerConfig` with the appropriate capabilities. /// /// The caller declares its own capability via `caller`. Sibling handler attributes /// (`prompt_handler`, `tool_handler`) are detected automatically @@ -153,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 1bdab33..a3e4ab7 100644 --- a/crates/rmcp/CHANGELOG.md +++ b/crates/rmcp/CHANGELOG.md @@ -7,6 +7,169 @@ 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 diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 93f98e4..ac698f9 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -18,6 +18,7 @@ exhaustive_enums = "warn" features = [ "auth", "auth-client-credentials-jwt", + "auth-enterprise-managed", "base64", "client", "client-side-sse", @@ -45,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" @@ -58,17 +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, features = ["aws_lc_rs"] } +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 = [ @@ -85,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 } @@ -127,7 +130,7 @@ 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", "base64"] +request-state = ["dep:hmac", "dep:sha2", "dep:zeroize", "base64"] # reqwest http client __reqwest = ["dep:reqwest"] @@ -187,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"] } @@ -210,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"] @@ -238,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"] @@ -288,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"] @@ -333,6 +348,15 @@ 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" @@ -385,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 = [ diff --git a/crates/rmcp/README.md b/crates/rmcp/README.md index 742b819..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/draft). 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 d61070b..486fa3b 100644 --- a/crates/rmcp/src/handler/client.rs +++ b/crates/rmcp/src/handler/client.rs @@ -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_subscriptions_acknowledged( - &self, - params: SubscriptionsAcknowledgedNotificationParams, - 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_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 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() } } @@ -400,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 39b89da..f846724 100644 --- a/crates/rmcp/src/handler/server.rs +++ b/crates/rmcp/src/handler/server.rs @@ -154,7 +154,8 @@ impl Service for H { McpError::method_not_found::(), ); }; - let advertised = requested.supported_by(&self.get_info().capabilities); + 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 { @@ -168,15 +169,16 @@ impl Service for H { "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 integrated draft schema defines a final result for graceful + // The 2026-07-28 schema defines a final result for graceful // server teardown; explicit stdio cancellation remains a notification. self.listen(subscription).await.map(|()| { - ServerResult::SubscriptionsListenResult( - SubscriptionsListenResult::complete(subscription_id), - ) + let mut result = SubscriptionsListenResult::complete(subscription_id); + result.meta.set_server_info(server_implementation); + ServerResult::SubscriptionsListenResult(result) }) } } @@ -298,6 +300,10 @@ 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 { @@ -315,14 +321,71 @@ 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) } @@ -405,7 +468,7 @@ macro_rules! server_handler_methods { /// /// The SDK sends the acknowledgment before invoking this method. Returning /// `Ok(())` sends the final [`SubscriptionsListenResult`] defined by the - /// integrated draft schema, marking graceful server teardown. Explicit + /// 2026-07-28 schema, marking graceful server teardown. Explicit /// stdio cancellation uses `notifications/cancelled` instead. fn listen( &self, @@ -537,8 +600,8 @@ macro_rules! server_handler_methods { std::future::ready(()) } - fn get_info(&self) -> ServerInfo { - ServerInfo::default() + fn get_info(&self) -> ServerConfig { + ServerConfig::default() } /// SEP-2663 `tasks/get`: return the current [`DetailedTask`] state. @@ -605,6 +668,13 @@ 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() } @@ -769,7 +839,7 @@ 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() } diff --git a/crates/rmcp/src/handler/server/router.rs b/crates/rmcp/src/handler/server/router.rs index e934137..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, }; @@ -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)] diff --git a/crates/rmcp/src/handler/server/tool.rs b/crates/rmcp/src/handler/server/tool.rs index a902406..54de5fa 100644 --- a/crates/rmcp/src/handler/server/tool.rs +++ b/crates/rmcp/src/handler/server/tool.rs @@ -34,10 +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, + /// 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> { @@ -47,6 +55,8 @@ impl<'s, S> ToolCallContext<'s, S> { meta: _, name, arguments, + input_responses, + request_state, .. }: CallToolRequestParams, request_context: RequestContext, @@ -56,6 +66,8 @@ impl<'s, S> ToolCallContext<'s, S> { service, name, arguments, + input_responses, + request_state, } } pub fn name(&self) -> &str { @@ -98,6 +110,12 @@ impl IntoCallToolResult for InputRequiredResult { } } +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 { Err(self) @@ -193,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 diff --git a/crates/rmcp/src/lib.rs b/crates/rmcp/src/lib.rs index 7c9b7b1..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; diff --git a/crates/rmcp/src/model.rs b/crates/rmcp/src/model.rs index 307ce52..ec7b93b 100644 --- a/crates/rmcp/src/model.rs +++ b/crates/rmcp/src/model.rs @@ -177,7 +177,7 @@ impl ProtocolVersion { /// 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. + /// All protocol versions known to this SDK, oldest first. pub const KNOWN_VERSIONS: &[Self] = &[ Self::V_2024_11_05, Self::V_2025_03_26, @@ -190,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 { @@ -503,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, @@ -1029,10 +1066,6 @@ impl RequestParamsMeta for InitializeRequestParams { } } -/// 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 @@ -1086,9 +1119,120 @@ 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`]. @@ -1121,8 +1265,8 @@ pub type DiscoverRequest = Request /// The server's response to a [`DiscoverRequest`]. #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -#[serde(rename_all = "camelCase")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[serde(rename_all = "camelCase")] #[non_exhaustive] pub struct DiscoverResult { /// Identifies how the result should be parsed. @@ -1131,8 +1275,6 @@ pub struct DiscoverResult { pub supported_versions: Vec, /// Capabilities provided by this server. pub capabilities: ServerCapabilities, - /// Information about the server implementation. - pub server_info: Implementation, /// Optional guidance for using the server. #[serde(skip_serializing_if = "Option::is_none")] pub instructions: Option, @@ -1145,18 +1287,26 @@ pub struct DiscoverResult { 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, - server_info: Implementation, - ) -> Self { + pub fn new(supported_versions: Vec, capabilities: ServerCapabilities) -> Self { Self { result_type: ResultType::COMPLETE, supported_versions, capabilities, - server_info, instructions: None, ttl_ms: 0, cache_scope: CacheScope::Private, @@ -1164,21 +1314,44 @@ impl DiscoverResult { } } - /// Create a discovery result from the server's initialization information. + /// 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_info: ServerInfo, + server_config: ServerConfig, ) -> Self { - let ServerInfo { + let ServerConfig { capabilities, server_info, instructions, meta, .. - } = server_info; - let mut result = Self::new(supported_versions, capabilities, server_info); - result.instructions = instructions; - result.meta = 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 } @@ -1195,10 +1368,24 @@ impl DiscoverResult { } } +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(), @@ -1209,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(), @@ -1390,9 +1577,6 @@ impl RequestParamsMeta for PaginatedRequestParams { } } -/// Deprecated: Use [`PaginatedRequestParams`] instead (SEP-1319 compliance). -#[deprecated(since = "0.13.0", note = "Use PaginatedRequestParams instead")] -pub type PaginatedRequestParam = PaginatedRequestParams; // ============================================================================= // PROGRESS AND PAGINATION // ============================================================================= @@ -1494,7 +1678,7 @@ macro_rules! paginated_result { /// the server handler clears the field when responding to peers that /// negotiated an older version. /// - /// [spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234 + /// [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")] @@ -1633,10 +1817,6 @@ impl RequestParamsMeta for ReadResourceRequestParams { } } -/// 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")] @@ -1652,7 +1832,7 @@ pub struct ReadResourceResult { /// the server handler clears the field when responding to peers that /// negotiated an older version. /// - /// [spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234 + /// [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). @@ -1741,10 +1921,6 @@ impl RequestParamsMeta for SubscribeRequestParams { } } -/// 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" @@ -1784,10 +1960,6 @@ impl RequestParamsMeta for UnsubscribeRequestParams { } } -/// 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" @@ -2018,7 +2190,7 @@ fn subscriptions_listen_request_meta_schema( #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[non_exhaustive] pub struct SubscriptionsListenRequestParams { - /// Protocol-level metadata. Required by the draft wire schema. + /// Protocol-level metadata. Required by the 2026-07-28 wire schema. #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] #[cfg_attr( feature = "schemars", @@ -2092,6 +2264,16 @@ impl SubscriptionsListenResultMeta { 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 { @@ -2130,9 +2312,14 @@ impl schemars::JsonSchema for 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"], @@ -2286,10 +2473,6 @@ impl RequestParamsMeta for GetPromptRequestParams { } } -/// 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; @@ -2359,10 +2542,6 @@ impl RequestParamsMeta for SetLevelRequestParams { } } -/// 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", @@ -2632,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 { @@ -2961,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 @@ -3154,10 +3326,6 @@ impl RequestParamsMeta for CompleteRequestParams { } } -/// 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)] @@ -3253,7 +3421,7 @@ pub struct CompleteResult { /// the server handler clears the field when responding to peers that /// negotiated an older version. /// - /// [spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234 + /// [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, @@ -3344,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] @@ -3498,21 +3663,20 @@ pub enum ElicitationAction { Cancel, } -/// Helper enum for deserializing CreateElicitationRequestParam with backward compatibility. -/// When mode is missing, it defaults to FormElicitationParam. +/// Wire representation for tagged elicitation parameters and legacy forms without `mode`. #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] #[serde(tag = "mode")] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -enum CreateElicitationRequestParamDeserializeHelper { +enum ElicitRequestParamsWire { #[serde(rename = "form", rename_all = "camelCase")] - FormElicitationParam { + Form { #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] 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, message: String, @@ -3520,7 +3684,7 @@ enum CreateElicitationRequestParamDeserializeHelper { elicitation_id: String, }, #[serde(untagged, rename_all = "camelCase")] - FormElicitationParamBackwardsCompat { + LegacyForm { #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] meta: Option, message: String, @@ -3528,19 +3692,17 @@ enum CreateElicitationRequestParamDeserializeHelper { }, } -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, @@ -3549,7 +3711,7 @@ impl TryFrom for ElicitRequestPa message, requested_schema, }), - CreateElicitationRequestParamDeserializeHelper::UrlElicitationParam { + ElicitRequestParamsWire::Url { meta, message, url, @@ -3595,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 { @@ -3650,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 @@ -3703,15 +3855,9 @@ impl ElicitResult { } } -#[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; - // ============================================================================= // TOOL EXECUTION RESULTS // ============================================================================= @@ -3734,7 +3880,7 @@ pub struct CallToolResult { /// the server handler clears the field when responding to peers that /// negotiated an older version. /// - /// [spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234 + /// [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.) @@ -3756,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 @@ -3775,6 +3922,16 @@ impl<'de> Deserialize<'de> for CallToolResult { 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() @@ -4041,10 +4198,6 @@ impl RequestParamsMeta for CallToolRequestParams { } } -/// 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; @@ -4122,7 +4275,7 @@ pub struct GetPromptResult { /// the server handler clears the field when responding to peers that /// negotiated an older version. /// - /// [spec schema]: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234 + /// [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")] @@ -4579,11 +4732,48 @@ mod tests { use super::*; #[test] - #[allow(deprecated)] - fn deprecated_aliases_still_resolve() { - // 하위호환: 구 이름이 새 타입으로 여전히 resolve되는지 확인. - let _: CreateElicitationResult = ElicitResult::new(ElicitationAction::Accept); - 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")] @@ -4806,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, @@ -5077,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": { @@ -5104,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"); } } @@ -5142,7 +5330,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"); } let json_data_url = json!({ @@ -5171,7 +5359,7 @@ mod tests { assert_eq!(url, "https://example.com/form"); assert_eq!(elicitation_id, "elicitation-123"); } else { - panic!("Expected UrlElicitationParam"); + panic!("Expected UrlElicitationParams"); } } diff --git a/crates/rmcp/src/model/elicitation_schema.rs b/crates/rmcp/src/model/elicitation_schema.rs index 29128fa..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 // ============================================================================= @@ -1112,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, @@ -1128,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>, @@ -1137,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, } @@ -1235,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() @@ -1600,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()); @@ -1673,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 { @@ -1862,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 0f7d90c..0a7121e 100644 --- a/crates/rmcp/src/model/meta.rs +++ b/crates/rmcp/src/model/meta.rs @@ -243,15 +243,6 @@ variant_extension! { #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] pub struct MetaObject(pub JsonObject); -/// Deprecated alias for [`MetaObject`]. -/// -/// This is a re-export rather than a type alias so the `Meta(...)` tuple -/// constructor keeps working. Request and notification metadata now have -/// dedicated types; use [`RequestMetaObject`] or [`NotificationMetaObject`] -/// where those are expected. -#[deprecated(note = "Use MetaObject (or RequestMetaObject / NotificationMetaObject)")] -pub use self::MetaObject as Meta; - impl MetaObject { /// Reserved `_meta` key for the W3C Trace Context `traceparent` value (SEP-414). const TRACEPARENT_FIELD: &str = "traceparent"; @@ -367,7 +358,7 @@ impl schemars::JsonSchema for MetaObject { fn json_schema(_: &mut schemars::SchemaGenerator) -> schemars::Schema { schemars::json_schema!({ - "description": "See [specification/draft/basic/index#general-fields] for notes on _meta usage.", + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", "type": "object", "additionalProperties": true, }) @@ -383,10 +374,10 @@ impl schemars::JsonSchema for MetaObject { /// - `io.modelcontextprotocol/clientCapabilities` (SEP-2575) /// - `io.modelcontextprotocol/logLevel` (SEP-2575) /// -/// The 2026-07-28 draft schema marks the protocol-version, client-info, and -/// client-capabilities keys as required; earlier protocol versions do not know -/// them. All keys therefore stay optional at runtime and in the generated -/// (version-shared) JSON schema — use +/// 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. /// @@ -406,9 +397,8 @@ impl RequestMetaObject { 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; 3] = [ + pub const DRAFT_REQUIRED_KEYS: [&str; 2] = [ Self::META_KEY_PROTOCOL_VERSION, - Self::META_KEY_CLIENT_INFO, Self::META_KEY_CLIENT_CAPABILITIES, ]; @@ -519,7 +509,7 @@ impl RequestMetaObject { /// meta.missing_required_keys(&ProtocolVersion::V_2025_11_25) /// .is_empty() /// ); - /// // The 2026-07-28 draft requires the SEP-2575 keys. + /// // 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(), @@ -533,9 +523,6 @@ impl RequestMetaObject { if self.protocol_version().is_none() { missing.push(Self::META_KEY_PROTOCOL_VERSION); } - if self.client_info().is_none() { - missing.push(Self::META_KEY_CLIENT_INFO); - } if self.client_capabilities().is_none() { missing.push(Self::META_KEY_CLIENT_CAPABILITIES); } @@ -586,9 +573,9 @@ impl schemars::JsonSchema for RequestMetaObject { 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 the 2026-07-28 draft marks as required are left + // 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. Draft-strict validation is available at + // `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.", @@ -848,7 +835,6 @@ mod tests { fn treats_malformed_values_as_missing() { let meta: RequestMetaObject = serde_json::from_value(serde_json::json!({ "io.modelcontextprotocol/protocolVersion": 123, - "io.modelcontextprotocol/clientInfo": "not an implementation", "io.modelcontextprotocol/clientCapabilities": null, })) .unwrap(); diff --git a/crates/rmcp/src/model/mrtr.rs b/crates/rmcp/src/model/mrtr.rs index 40621a7..8dac9a3 100644 --- a/crates/rmcp/src/model/mrtr.rs +++ b/crates/rmcp/src/model/mrtr.rs @@ -274,6 +274,12 @@ impl<'de> Deserialize<'de> for InputRequiredResult { } } + 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, @@ -451,6 +457,19 @@ mod tests { ); } + #[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!({ diff --git a/crates/rmcp/src/model/request_state.rs b/crates/rmcp/src/model/request_state.rs index 5ed0f7d..9f45572 100644 --- a/crates/rmcp/src/model/request_state.rs +++ b/crates/rmcp/src/model/request_state.rs @@ -40,9 +40,10 @@ //! //! ``` //! 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::new(b"a-32-byte-or-longer-secret-key!!!"); +//! 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"; @@ -53,38 +54,60 @@ //! //! // 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).expect("integrity check passes"); +//! 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::time::Duration; +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; -/// Version tag prefixing every sealed value, so the wire format can evolve. -const VERSION: &str = "rs1"; +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: &[u8] = b"rmcp/mrtr/request-state/v1"; +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; -/// Errors returned when opening a sealed [`RequestStateCodec`] value. +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")] @@ -103,6 +126,18 @@ pub enum RequestStateError { #[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), @@ -148,35 +183,261 @@ impl<'a> SealOptions<'a> { } } -/// A keyed codec that seals and opens SEP-2322 `requestState` values with -/// HMAC-SHA256 integrity protection. +/// A keyed codec for integrity-protected SEP-2322 `requestState` values. /// -/// Construct one codec per signing key and reuse it for the lifetime of the -/// key. The same key must be used to [`seal`](Self::seal) and -/// [`open`](Self::open) a value, so it has to survive across the rounds of a -/// single MRTR exchange (e.g. a stable per-process or per-deployment secret). +/// [`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. /// -/// The key may be any length; HMAC internally normalizes it. For meaningful -/// security use a high-entropy key of at least 32 bytes. +/// 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 { - key: Box<[u8]>, + 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 the signing key through Debug output. - f.debug_struct("RequestStateCodec") - .field("key", &"") - .finish() + // 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 { - /// Creates a codec from a signing key. + /// 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 { - key: key.into().into_boxed_slice(), + 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) + } } } @@ -235,12 +496,12 @@ impl RequestStateCodec { /// /// # Errors /// - /// - [`RequestStateError::IntegrityCheckFailed`] if the value was not - /// produced by this key or the associated data differs. - /// - [`RequestStateError::Expired`] if the value's TTL has elapsed. - /// - [`RequestStateError::MalformedFormat`] or - /// [`RequestStateError::InvalidEncoding`] if it is not a well-formed sealed - /// value. + /// 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, @@ -287,24 +548,76 @@ impl RequestStateCodec { body.extend_from_slice(&expiry.to_be_bytes()); body.extend_from_slice(payload); - let tag = self - .mac_for(options.associated_data, &body) + 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(); - - // base64url without padding encodes 3 bytes as 4 chars, rounding up. - let b64_len = |n: usize| n.div_ceil(3) * 4; - let mut out = - String::with_capacity(VERSION.len() + 2 + b64_len(body.len()) + b64_len(tag.len())); - out.push_str(VERSION); + 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); + 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_at( + 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], @@ -314,10 +627,17 @@ impl RequestStateCodec { 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 { + 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)?; @@ -325,11 +645,81 @@ impl RequestStateCodec { .decode(tag_b64) .map_err(|_| RequestStateError::InvalidEncoding)?; - // `verify_slice` compares in constant time and rejects wrong-length tags. - self.mac_for(associated_data, &body) + // 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); @@ -342,20 +732,55 @@ impl RequestStateCodec { Ok(body[EXPIRY_LEN..].to_vec()) } - /// Builds an HMAC keyed for request-state tags, pre-fed with the - /// domain-separation label, a length-prefixed `associated_data`, and the - /// body. The length prefix keeps the `associated_data`/`body` boundary - /// unambiguous so distinct inputs cannot collide. - fn mac_for(&self, associated_data: &[u8], body: &[u8]) -> HmacSha256 { - let mut mac = - HmacSha256::new_from_slice(&self.key).expect("HMAC accepts keys of any length"); - mac.update(DOMAIN); + 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() } @@ -365,9 +790,42 @@ impl RequestStateCodec { 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::new(b"test-key-test-key-test-key-32byte".to_vec()); + 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"); @@ -380,7 +838,8 @@ mod tests { tool: String, round: u32, } - let codec = RequestStateCodec::new(b"another-strong-signing-key-here!!".to_vec()); + let codec = + RequestStateCodec::try_new(b"another-strong-signing-key-here!!".to_vec()).unwrap(); let state = State { tool: "weather".into(), round: 3, @@ -392,14 +851,16 @@ mod tests { #[test] fn empty_payload_roundtrips() { - let codec = RequestStateCodec::new(b"k".to_vec()); + 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::new(b"signing-key-signing-key-signing!!".to_vec()); + 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. @@ -416,8 +877,10 @@ mod tests { #[test] fn different_key_is_rejected() { - let signer = RequestStateCodec::new(b"the-real-signing-key-value-here!!".to_vec()); - let attacker = RequestStateCodec::new(b"a-totally-different-forged-key!!!".to_vec()); + 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), @@ -427,17 +890,19 @@ mod tests { #[test] fn appended_bytes_are_rejected() { - let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec()); + 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 wrong_version_prefix_is_malformed() { - let codec = RequestStateCodec::new(b"key".to_vec()); + 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.", "rs2.", 1); + let bumped = sealed.replacen("rs1.", "rs3.", 1); assert!(matches!( codec.open(&bumped), Err(RequestStateError::MalformedFormat) @@ -446,7 +911,8 @@ mod tests { #[test] fn missing_sections_are_malformed() { - let codec = RequestStateCodec::new(b"key".to_vec()); + let codec = + RequestStateCodec::try_new(b"missing-sections-test-signing-key".to_vec()).unwrap(); assert!(matches!( codec.open("rs1"), Err(RequestStateError::MalformedFormat) @@ -463,7 +929,8 @@ mod tests { #[test] fn non_base64_sections_are_invalid_encoding() { - let codec = RequestStateCodec::new(b"key".to_vec()); + let codec = + RequestStateCodec::try_new(b"invalid-base64-test-signing-key!!".to_vec()).unwrap(); assert!(matches!( codec.open("rs1.!!!!.!!!!"), Err(RequestStateError::InvalidEncoding) @@ -472,7 +939,8 @@ mod tests { #[test] fn debug_does_not_leak_key() { - let codec = RequestStateCodec::new(b"super-secret-key".to_vec()); + 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")); @@ -483,7 +951,8 @@ mod tests { #[test] fn matching_context_opens() { - let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec()); + 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"); @@ -491,7 +960,8 @@ mod tests { #[test] fn different_context_is_rejected() { - let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec()); + 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!( @@ -502,7 +972,8 @@ mod tests { #[test] fn missing_context_is_rejected() { - let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec()); + 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. @@ -520,7 +991,7 @@ mod tests { #[test] fn within_ttl_opens() { - let codec = RequestStateCodec::new(KEY.to_vec()); + let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); let sealed = codec.seal_at( b"state", &SealOptions::new().ttl(Duration::from_secs(60)), @@ -532,7 +1003,7 @@ mod tests { #[test] fn past_ttl_is_expired() { - let codec = RequestStateCodec::new(KEY.to_vec()); + let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); let sealed = codec.seal_at( b"state", &SealOptions::new().ttl(Duration::from_secs(60)), @@ -547,14 +1018,14 @@ mod tests { #[test] fn no_ttl_never_expires() { - let codec = RequestStateCodec::new(KEY.to_vec()); + 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::new(KEY.to_vec()); + let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); let ctx = b"user:alice"; let sealed = codec.seal_at( b"state", @@ -574,4 +1045,450 @@ mod tests { )); } } + + 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/serde_impl.rs b/crates/rmcp/src/model/serde_impl.rs index b974722..bfd3b57 100644 --- a/crates/rmcp/src/model/serde_impl.rs +++ b/crates/rmcp/src/model/serde_impl.rs @@ -86,9 +86,8 @@ struct ProxyNoParam { } /// Combine the message-specific `_meta` map with a legacy [`MetaObject`] -/// extension (inserted through the deprecated `Meta` name), so pre-3.x code -/// does not silently lose metadata on the wire. On key conflicts the -/// message-specific map wins. +/// 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, @@ -99,7 +98,11 @@ fn merge_legacy_meta<'a>( (None, Some(legacy)) => Some(Cow::Borrowed(legacy)), (Some(typed), Some(legacy)) => { let mut merged = legacy.clone(); - merged.extend(typed.clone()); + merged.extend( + typed + .iter() + .map(|(key, value)| (key.clone(), value.clone())), + ); Some(Cow::Owned(merged)) } (None, None) => None, @@ -778,8 +781,6 @@ mod test { #[test] fn test_legacy_meta_extension_still_serializes() { - // Pre-3.x code inserts `MetaObject` into extensions through the - // deprecated `Meta` name; its metadata must not be silently dropped. let mut extensions = Extensions::new(); let mut legacy = crate::model::MetaObject::new(); legacy.insert("traceId".to_string(), json!("legacy")); diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs index b1b2c2b..63b38e4 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -1,4 +1,4 @@ -use std::sync::OnceLock; +use std::{borrow::Cow, sync::OnceLock}; use futures::FutureExt; #[cfg(not(feature = "local"))] @@ -136,6 +136,14 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone { 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 @@ -168,18 +176,61 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone { fn enforce_peer_request_association( _peer_request: &Self::PeerReq, _peer_info: Option<&Self::PeerInfo>, - _has_pending_outbound_request: bool, + _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(|version| version < &ProtocolVersion::V_2026_07_28) + !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! { @@ -205,10 +256,26 @@ pub(crate) fn in_request_handler_scope() -> bool { /// 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 = JsonRpcMessage<::Req, ::Resp, ::Not>; pub type RxJsonRpcMessage = JsonRpcMessage< @@ -230,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")] @@ -245,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 { @@ -296,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"))] @@ -311,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")] @@ -326,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 { @@ -346,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::{ @@ -429,6 +530,10 @@ 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>; @@ -1205,7 +1310,14 @@ where { let (peer, peer_rx) = Peer::new(Arc::new(AtomicU32RequestIdProvider::default()), peer_info); R::configure_direct_peer(&peer, &service.get_info()); - serve_inner(service, transport.into_transport(), peer, peer_rx, ct) + 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. @@ -1230,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 @@ -1255,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 @@ -1268,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)] @@ -1352,7 +1471,7 @@ where Event::SendTaskResult(SendTaskResult::Request { id, result }) => { if let Err(e) = result && let Some(responder) = local_responder_pool.remove(&id) { - let _ = responder.send(Err(ServiceError::TransportSend(e))); + let _ = responder.responder.send(Err(ServiceError::TransportSend(e))); } } Event::SendTaskResult(SendTaskResult::Notification { @@ -1370,7 +1489,7 @@ where && 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(), })); } @@ -1407,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(); @@ -1448,7 +1570,7 @@ where if let Err(error) = R::enforce_peer_request_association( &request, peer.peer_info().as_deref(), - !local_responder_pool.is_empty(), + 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 @@ -1508,31 +1630,38 @@ where })) => { tracing::info!(?notification, "received notification"); R::invalidate_response_cache(&peer, ¬ification).await; - let cancellation_request_id = + let subscription_id = if let Some(cancelled) = R::peer_cancelled_params(¬ification) { let request_id = cancelled.request_id.clone(); - if let Some(request_id) = request_id.as_ref() { - if R::IS_CLIENT { - if let Some(responder) = - local_responder_pool.remove(request_id) + // 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 _ = responder.send(Err(ServiceError::Cancelled { + let _ = pending.responder.send(Err(ServiceError::Cancelled { reason: cancelled.reason.clone(), })); + } else { + tracing::debug!(%request_id, "ignoring cancellation of unknown subscription"); + continue; } - } else if let Some(ct) = local_ct_pool.remove(request_id) { + } + 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 } - request_id } else { - None + notification.get_meta().subscription_id() }; - let subscription_id = notification - .get_meta() - .subscription_id() - .or(cancellation_request_id); if let Some(subscription_id) = subscription_id && let Some((sender, capacity)) = peer.subscription_sender(&subscription_id) @@ -1549,7 +1678,7 @@ where && let Some(responder) = local_responder_pool.remove(&subscription_id) { - let _ = responder + let _ = responder.responder .send(Err(ServiceError::SubscriptionLagged { capacity })); } peer.unregister_subscription(&subscription_id); @@ -1601,7 +1730,7 @@ where if let Some(responder) = remove_pending_request(&mut local_responder_pool, &id) { - let response_result = responder.send(Ok(result)); + let response_result = responder.responder.send(Ok(result)); if let Err(_error) = response_result { tracing::warn!(%id, "Error sending response"); } @@ -1622,7 +1751,7 @@ where } else { ServiceError::McpError(error) }; - let _response_result = responder.send(Err(service_error)); + let _response_result = responder.responder.send(Err(service_error)); if let Err(_error) = _response_result { tracing::warn!(%id, "Error sending response"); } @@ -1656,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; @@ -1729,4 +1869,73 @@ mod sep2260_marker_tests { 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 40d410d..7d0d794 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -13,19 +13,19 @@ use crate::{ model::{ ArgumentInfo, CacheScope, CallToolRequest, CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, CancelTaskRequest, CancelledNotification, - CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage, ClientNotification, + 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, InitializedNotification, InputRequest, - InputRequiredResult, InputResponses, JsonRpcResponse, ListPromptsRequest, + 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, - ServerInfo, ServerJsonRpcMessage, ServerNotification, ServerRequest, ServerResult, + ServerJsonRpcMessage, ServerNotification, ServerPeerInfo, ServerRequest, ServerResult, SetLevelRequest, SetLevelRequestParams, SubscribeRequest, SubscribeRequestParams, SubscriptionFilter, SubscriptionsListenRequest, SubscriptionsListenRequestParams, SubscriptionsListenResult, UnsubscribeRequest, UnsubscribeRequestParams, UpdateTaskParams, @@ -49,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), @@ -74,6 +82,12 @@ pub enum ClientInitializeError { #[error("Cancelled")] Cancelled, + + #[error("discover and legacy initialize both failed")] + LegacyFallbackFailed { + discover: Box, + fallback: Box, + }, } impl ClientInitializeError { @@ -86,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 @@ -102,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, @@ -116,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) => { @@ -187,8 +266,8 @@ 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; @@ -213,14 +292,27 @@ impl ServiceRole for RoleClient { } } - // SEP-2260: with no outbound request in flight there is nothing the - // server request could be associated with, so reject it. With one in - // flight we cannot tell which request it belongs to (no wire field), so - // we accept — an under-approximation of the spec's SHOULD. + 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>, - has_pending_outbound_request: bool, + association: PeerRequestAssociation, ) -> Result<(), ErrorData> { let restricted = matches!( peer_request, @@ -233,13 +325,24 @@ impl ServiceRole for RoleClient { } let strict = peer_info.is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28); - if strict && !has_pending_outbound_request { - return Err(ErrorData::invalid_params( + 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, - )); + )) } - Ok(()) } async fn invalidate_response_cache(peer: &Peer, notification: &Self::PeerNot) { @@ -543,13 +646,16 @@ pub enum ClientLifecycleMode { Discover { preferred_versions: Vec, }, - /// Probe with `server/discover`, falling back only when the peer proves it is legacy. + /// 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( @@ -641,7 +747,13 @@ where E: std::error::Error + Send + Sync + 'static, { tokio::select! { - result = serve_client_with_ct_inner(service, transport.into_transport(), lifecycle, 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) } @@ -653,6 +765,7 @@ async fn serve_client_with_ct_inner( transport: T, lifecycle: ClientLifecycleMode, ct: CancellationToken, + auto_discover_timeout: Duration, ) -> Result, ClientInitializeError> where S: Service, @@ -668,7 +781,7 @@ where legacy_startup(&service, &mut transport, &id_provider, &peer, client_info).await?; } ClientLifecycleMode::Discover { preferred_versions } => { - discover_startup( + match discover_startup( &service, &mut transport, &id_provider, @@ -676,26 +789,48 @@ where &client_info, preferred_versions, ) - .await?; + .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 = discover_startup( - &service, - &mut transport, - &id_provider, - &peer, - &client_info, - preferred_versions, + 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(()) => {} - Err(ClientInitializeError::JsonRpcError(error)) - if error.code == crate::model::ErrorCode::METHOD_NOT_FOUND => - { + 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; @@ -703,11 +838,52 @@ where legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info) .await?; } - Err(error) => return Err(error), } } } - Ok(serve_inner(service, transport, peer, peer_rx, ct)) + 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( @@ -715,7 +891,7 @@ async fn legacy_startup( transport: &mut T, id_provider: &Arc, peer: &Peer, - client_info: ClientInfo, + client_info: InitializeRequestParams, ) -> Result<(), ClientInitializeError> where S: Service, @@ -738,20 +914,13 @@ where context: "send initialize request".into(), })?; - let (response, response_id) = - expect_response(transport, "initialize response", service, peer.clone()).await?; - - if !id.matches_response_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( @@ -771,9 +940,9 @@ async fn discover_startup( transport: &mut T, id_provider: &Arc, peer: &Peer, - client_info: &ClientInfo, + client_info: &InitializeRequestParams, preferred_versions: Vec, -) -> Result<(), ClientInitializeError> +) -> Result where S: Service, T: Transport + 'static, @@ -805,14 +974,8 @@ where ClientInitializeError::transport::(error, "send discover request") })?; - match expect_response(transport, "discover response", service, peer.clone()).await { - Ok((ServerResult::DiscoverResult(result), response_id)) => { - if !id.matches_response_id(&response_id) { - return Err(ClientInitializeError::ConflictInitResponseId( - id, - response_id, - )); - } + 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 { @@ -821,21 +984,18 @@ where server_supported: result.supported_versions, }); }; - peer.set_peer_info(ServerInfo { - protocol_version: selected.clone(), - capabilities: result.capabilities, - server_info: result.server_info, - instructions: result.instructions, - meta: result.meta, - }); + 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(()); + return Ok(DiscoverOutcome::Modern); } - Ok((response, _)) => { + Ok(response) => { return Err(ClientInitializeError::ExpectedInitResult(Some(response))); } Err(ClientInitializeError::JsonRpcError(error)) @@ -869,6 +1029,19 @@ where }; 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), } } @@ -2041,6 +2214,300 @@ where 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); @@ -2172,13 +2639,10 @@ mod tests { let peer = disconnected_peer(); let meta = RequestMetaObject::default(); let key = discover_cache_key(); - let expected = DiscoverResult::new( - vec![ProtocolVersion::default()], - Default::default(), - crate::model::Implementation::from_build_env(), - ) - .with_ttl_ms(5_000) - .with_cache_scope(CacheScope::Public); + 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()), @@ -2190,3 +2654,73 @@ mod tests { 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/server.rs b/crates/rmcp/src/service/server.rs index 2b84a94..3ff229f 100644 --- a/crates/rmcp/src/service/server.rs +++ b/crates/rmcp/src/service/server.rs @@ -13,15 +13,16 @@ use super::*; 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, SubscriptionFilter, SubscriptionsAcknowledgedNotification, - SubscriptionsAcknowledgedNotificationParams, 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, }; @@ -37,8 +38,8 @@ 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; @@ -85,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), @@ -101,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, @@ -474,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") + } } } @@ -535,33 +594,40 @@ where let initialize_request = match request { ClientRequest::InitializeRequest(request) => request, - mut request => { - if !request + request => { + let missing_metadata = request .get_meta() - .missing_required_keys(&ProtocolVersion::V_2026_07_28) - .is_empty() - { + .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(); - let context = RequestContext { - ct: ct.child_token(), - id: id.clone(), - meta: std::mem::take(request.get_meta_mut()), - extensions: std::mem::take(request.extensions_mut()), - peer: peer.clone(), - }; - let response = match service.handle_request(request, context).await { - Ok(result) => ServerJsonRpcMessage::response(result, id), - Err(error) => ServerJsonRpcMessage::error(error, Some(id)), - }; - transport.send(response).await.map_err(|error| { - ServerInitializeError::transport::(error, "sending negotiated request response") - })?; - return Ok(serve_inner(service, transport, peer, peer_rx, ct)); + // 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 requested_protocol_version = initialize_request.params.protocol_version.clone(); @@ -576,24 +642,21 @@ where peer: peer.clone(), }; // Send initialize response - let init_response = service.handle_request(request, 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 = 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), }; - init_response.protocol_version = - negotiate_protocol_version(&requested_protocol_version, init_response.protocol_version); // Update peer_info so context.protocol_version() reflects the negotiated // version in all subsequent request handlers. negotiated_peer_info.protocol_version = init_response.protocol_version.clone(); @@ -614,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 { diff --git a/crates/rmcp/src/transport.rs b/crates/rmcp/src/transport.rs index 74a1394..fdedc9c 100644 --- a/crates/rmcp/src/transport.rs +++ b/crates/rmcp/src/transport.rs @@ -100,7 +100,7 @@ pub use auth::JwtSigningAlgorithm; #[cfg(feature = "auth")] pub use auth::{ AuthClient, AuthError, AuthorizationManager, AuthorizationRequest, AuthorizationSession, - AuthorizedHttpClient, ClientCredentialsConfig, CredentialStore, + AuthorizedHttpClient, ClientCredentialsConfig, CredentialRefreshGuard, CredentialStore, EXTENSION_OAUTH_CLIENT_CREDENTIALS, InMemoryCredentialStore, InMemoryStateStore, OAuthHttpClient, OAuthHttpClientError, OAuthHttpClientFuture, OAuthHttpRedirectPolicy, OAuthHttpRequest, ScopeUpgradeConfig, StateStore, StoredAuthorizationState, StoredCredentials, @@ -272,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 46d5dea..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" ) } @@ -582,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")); @@ -590,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 26e5378..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, @@ -28,15 +28,12 @@ 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; -const RESOURCE_METADATA_POST_PROBE_BODY: &str = concat!( - r#"{"jsonrpc":"2.0","id":"auth-discovery","method":"initialize","params":{"#, - r#""protocolVersion":"2024-11-05","capabilities":{},"#, - r#""clientInfo":{"name":"rmcp-auth-discovery","version":"0.0.0"}}"#, - r#"}"# -); const CLOUD_METADATA_HOSTS: &[&str] = &[ "metadata", "metadata.google.internal", @@ -76,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`]. @@ -106,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, @@ -138,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()) @@ -153,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) }) } } @@ -173,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) + }) } } @@ -243,11 +271,32 @@ impl StoredCredentials { } } +/// 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), + } + } +} + /// Trait for storing and retrieving OAuth2 credentials /// /// 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>; @@ -255,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) @@ -503,6 +562,9 @@ pub enum AuthError { #[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), @@ -595,7 +657,7 @@ pub enum AuthorizationMetadataSource { /// [Newer MCP revisions] require metadata discovery and do not define an /// endpoint-synthesis fallback. /// - /// [Newer MCP revisions]: https://modelcontextprotocol.io/specification/draft/basic/authorization/authorization-server-discovery#protected-resource-metadata-discovery-requirements + /// [Newer MCP revisions]: https://modelcontextprotocol.io/specification/2026-07-28/basic/authorization/authorization-server-discovery#protected-resource-metadata-discovery-requirements LegacyEndpointFallback, } @@ -623,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] @@ -634,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") @@ -687,7 +768,7 @@ impl OAuthClientConfig { /// Declarative description of the client identity material available for an /// authorization flow. /// -/// The [MCP authorization specification](https://modelcontextprotocol.io/specification/draft/basic/authorization/client-registration) +/// 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: @@ -736,6 +817,11 @@ pub struct AuthorizationRequest { /// 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 { @@ -750,6 +836,7 @@ impl AuthorizationRequest { client_secret: None, client_metadata_url: None, application_type: None, + challenge: None, } } @@ -805,6 +892,13 @@ impl AuthorizationRequest { 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 @@ -883,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) @@ -906,6 +1000,42 @@ pub enum ClientCredentialsConfig { }, } +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, @@ -998,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)] @@ -1204,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 @@ -1244,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) @@ -1255,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, @@ -1392,7 +1534,7 @@ impl AuthorizationManager { }); } - if let Some(metadata) = self.try_discover_oauth_server(&self.base_url).await? { + if let Some(metadata) = self.try_discover_oauth_server(&self.base_url, None).await? { return Ok(AuthorizationMetadataResolution { metadata, source: AuthorizationMetadataSource::AuthorizationServerMetadata, @@ -1406,6 +1548,52 @@ impl AuthorizationManager { }) } + /// 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(); @@ -1655,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 { @@ -1693,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() @@ -1917,11 +2113,13 @@ impl AuthorizationManager { AuthError::InternalError("Authorization state not found".to_string()) })?; - // Delete state after retrieval (one-time use) - self.state_store.delete(csrf_token).await?; - 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?; + // capture requested scopes before the state is consumed let requested_scopes = stored_state.requested_scopes.clone(); @@ -1934,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, @@ -2043,7 +2241,7 @@ 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."); @@ -2064,8 +2262,21 @@ 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)?; @@ -2079,9 +2290,10 @@ impl AuthorizationManager { 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.base_url.to_string()); + .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)); } @@ -2107,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(); @@ -2197,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)); } } @@ -2209,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()); @@ -2226,7 +2443,11 @@ impl AuthorizationManager { match serde_json::from_slice::(response.body()) { Ok(metadata) => { - Self::validate_authorization_metadata_issuer(discovery_url, &metadata)?; + self.validate_authorization_metadata_issuer( + discovery_url, + expected_issuer, + &metadata, + )?; Ok(Some(metadata)) } Err(err) => { @@ -2288,15 +2509,21 @@ impl AuthorizationManager { } fn validate_authorization_metadata_issuer( + &self, discovery_url: &Url, + expected_issuer: Option<&str>, metadata: &AuthorizationMetadata, ) -> Result<(), AuthError> { - let Some(expected_issuer) = - Self::expected_issuer_for_authorization_metadata_url(discovery_url) - else { + 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) { @@ -2311,18 +2538,54 @@ impl AuthorizationManager { 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 @@ -2331,21 +2594,32 @@ impl AuthorizationManager { *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) { @@ -2363,18 +2637,31 @@ impl AuthorizationManager { } 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( @@ -2387,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 ))); } @@ -2397,129 +2696,140 @@ 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('/')) - || Self::root_resource_identifier_covers_path(actual, expected) - } - - 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()) - } - - fn root_resource_identifier_covers_path(root_resource: &str, path_resource: &str) -> bool { - let Ok(root_resource) = Url::parse(root_resource) else { - return false; - }; - let Ok(path_resource) = Url::parse(path_resource) else { - return false; - }; - - root_resource.path() == "/" - && root_resource.query().is_none() - && root_resource.fragment().is_none() - && path_resource.path() != "/" - && Self::is_same_origin(&root_resource, &path_resource) - } - - async fn discover_resource_metadata_url(&self) -> Result, AuthError> { - if let Ok(Some(resource_metadata_url)) = - self.fetch_resource_metadata_url(&self.base_url, true).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, false) + 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 - { - return Ok(Some(resource_metadata_url)); + .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, - allow_post_probe: bool, - ) -> Result, AuthError> { - let response = match self.discovery_get(url).await { - Ok(r) => r, - Err(e) => { - debug!("resource metadata probe failed: {}", e); - return Ok(None); - } - }; - - match response.status() { - StatusCode::OK => Ok(Some(url.clone())), - StatusCode::UNAUTHORIZED => Ok(self - .extract_resource_metadata_url_from_www_authenticate(&response) - .await), - StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED if allow_post_probe => { - self.fetch_resource_metadata_url_with_post_probe(url).await - } - status => { - debug!("resource metadata probe returned unexpected status: {status}"); - Ok(None) - } - } - } - - async fn fetch_resource_metadata_url_with_post_probe( - &self, - url: &Url, - ) -> Result, AuthError> { - let request = oauth2::http::Request::builder() - .method("POST") - .uri(url.as_str()) - .header(HEADER_MCP_PROTOCOL_VERSION, "2024-11-05") - .header(CONTENT_TYPE, "application/json") - .body(RESOURCE_METADATA_POST_PROBE_BODY.as_bytes().to_vec()) - .map_err(|error| AuthError::InternalError(error.to_string()))?; - let response = match self - .http_client - .execute(OAuthHttpRequest::new( - request, - OAuthHttpRedirectPolicy::Stop, - )) + async fn probe_resource_endpoint_for_challenge(&self) -> Result, AuthError> { + let response = self + .discovery_get(&self.base_url) .await - { - Ok(response) => response, - Err(error) => { - debug!("resource metadata POST probe failed: {}", error); - return Ok(None); - } - }; + .map_err(|error| Self::discovery_failed(&self.base_url, error))?; - if response.status() != StatusCode::UNAUTHORIZED { - debug!( - "resource metadata POST 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); } - 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( @@ -2532,14 +2842,9 @@ impl AuthorizationManager { 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; } } @@ -2550,18 +2855,16 @@ impl AuthorizationManager { 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!( @@ -2571,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 { @@ -2589,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, )) @@ -2607,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, ))) } @@ -3156,7 +3536,7 @@ pub struct AuthorizationSession { impl AuthorizationSession { /// Create a new authorization session, selecting a client registration - /// mechanism per the [MCP authorization specification](https://modelcontextprotocol.io/specification/draft/basic/authorization/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 @@ -3415,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( @@ -3511,7 +3891,7 @@ impl OAuthState { /// Start authorization. /// /// Selects a client registration mechanism from the identity material in - /// `request`, following the [MCP authorization specification](https://modelcontextprotocol.io/specification/draft/basic/authorization/client-registration) + /// `request`, following the [MCP authorization specification](https://modelcontextprotocol.io/specification/2026-07-28/basic/authorization/client-registration) /// priority order: /// /// 1. Pre-registered client information @@ -3531,7 +3911,7 @@ impl OAuthState { &mut self, request: AuthorizationRequest, ) -> Result<(), AuthError> { - let placeholder = self.placeholder().await?; + let placeholder = self.placeholder_state().await?; let old = std::mem::replace(self, placeholder); let OAuthState::Unauthorized(mut manager) = old else { *self = old; @@ -3540,7 +3920,10 @@ impl OAuthState { )); }; debug!("start discovery"); - let metadata = match manager.resolve_metadata().await { + 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); @@ -3563,7 +3946,7 @@ impl OAuthState { /// 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(()) @@ -3571,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), @@ -3593,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; @@ -3725,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(), @@ -3757,16 +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, - AuthorizationMetadataSource, AuthorizationRequest, AuthorizationSession, CredentialStore, - InMemoryCredentialStore, 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; @@ -3806,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()) }) } } @@ -3834,6 +4218,219 @@ mod tests { .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() @@ -3849,7 +4446,7 @@ 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"] }), ), @@ -3872,6 +4469,10 @@ mod tests { 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![ @@ -3899,16 +4500,113 @@ mod tests { } #[tokio::test] - async fn protected_resource_metadata_supports_authorization_server_path_insertion() { + 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![ - empty_response(401), + challenge, http_response( 200, serde_json::json!({ - "resource": "https://mcp.example.com/", - "authorization_servers": ["https://auth.example.com/tenant1"] + "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!({ @@ -3950,13 +4648,53 @@ mod tests { vec![ "https://mcp.example.com/", "https://mcp.example.com/.well-known/oauth-protected-resource", - "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![ @@ -4006,8 +4744,8 @@ mod tests { ); } - #[test] - fn authorization_metadata_accepts_oidc_path_appended_issuer() { + #[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(); @@ -4017,8 +4755,12 @@ mod tests { token_endpoint: "https://auth.example.com/tenant1/token".to_string(), ..Default::default() }; + let manager = AuthorizationManager::new("https://mcp.example.com/") + .await + .unwrap(); - AuthorizationManager::validate_authorization_metadata_issuer(&discovery_url, &metadata) + manager + .validate_authorization_metadata_issuer(&discovery_url, None, &metadata) .unwrap(); } @@ -4034,8 +4776,8 @@ mod tests { )); } - #[test] - fn authorization_metadata_accepts_oidc_path_inserted_issuer() { + #[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(); @@ -4045,13 +4787,17 @@ mod tests { token_endpoint: "https://auth.example.com/tenant1/token".to_string(), ..Default::default() }; + let manager = AuthorizationManager::new("https://mcp.example.com/") + .await + .unwrap(); - AuthorizationManager::validate_authorization_metadata_issuer(&discovery_url, &metadata) + manager + .validate_authorization_metadata_issuer(&discovery_url, None, &metadata) .unwrap(); } - #[test] - fn authorization_metadata_rejects_missing_issuer_for_standard_discovery_url() { + #[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(); @@ -4062,9 +4808,34 @@ mod tests { ..Default::default() }; - let error = - AuthorizationManager::validate_authorization_metadata_issuer(&discovery_url, &metadata) - .unwrap_err(); + 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!( @@ -4087,7 +4858,6 @@ mod tests { .body(Vec::new()) .unwrap(); let client = RecordingOAuthHttpClient::with_responses(vec![ - empty_response(404), challenge, http_response( 200, @@ -4128,7 +4898,6 @@ mod tests { ( "https://auth.example.com/tenant1/token", vec![ - "https://mcp.example.com/mcp", "https://mcp.example.com/mcp", "https://mcp.example.com/custom/metadata/location.json", "https://auth.example.com/.well-known/oauth-authorization-server/tenant1", @@ -4137,25 +4906,27 @@ mod tests { ], ) ); - assert_eq!( + assert!( client .requests() .iter() - .take(2) - .map(|request| request.method.as_str()) - .collect::>(), - vec!["GET", "POST"] + .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() { + async fn resolve_metadata_reports_legacy_fallback_when_nothing_is_discovered( + #[case] status: StatusCode, + ) { let client = RecordingOAuthHttpClient::with_responses(vec![ - empty_response(404), - empty_response(404), - empty_response(404), - empty_response(404), - empty_response(404), + 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/", @@ -4184,7 +4955,6 @@ mod tests { "https://legacy.example.com/token", Some("https://legacy.example.com/register"), vec![ - "https://legacy.example.com/", "https://legacy.example.com/", "https://legacy.example.com/.well-known/oauth-protected-resource", "https://legacy.example.com/.well-known/oauth-authorization-server", @@ -4243,6 +5013,85 @@ mod tests { ); } + #[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![ @@ -4279,6 +5128,579 @@ mod tests { ); } + #[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, @@ -4373,6 +5795,32 @@ mod tests { 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()); @@ -4988,35 +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::resource_identifiers_match( - "https://mcp.example.com/mcp", - "https://mcp.example.com" + 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::resource_identifiers_match( - "https://mcp.example.com/mcp", - "https://real.example.com" + 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::resource_identifiers_match( - "https://mcp.example.com/mcp", - "https://mcp.example.com?resource=mcp" + 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() )); } @@ -5976,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( @@ -6146,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] @@ -6774,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; @@ -7545,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/auth/streamable_http_client.rs b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs index 2069e6d..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,16 +81,21 @@ 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( @@ -33,18 +103,24 @@ where uri: 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( @@ -52,26 +128,32 @@ where uri: std::sync::Arc, session_id: Option>, last_event_id: Option, - mut auth_token: Option, + auth_token: Option, custom_headers: HashMap, max_sse_event_size: usize, ) -> 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_with_max_sse_event_size( - uri, - session_id, - last_event_id, - auth_token, - custom_headers, - max_sse_event_size, - ) - .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_with_max_sse_event_size( + uri, + session_id, + last_event_id, + token, + custom_headers, + max_sse_event_size, + ) + .await + } + }) + .await } async fn post_message( @@ -79,18 +161,24 @@ 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.http_client - .post_message(uri, message, session_id, auth_token, custom_headers) - .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( @@ -98,25 +186,195 @@ where uri: std::sync::Arc, message: crate::model::ClientJsonRpcMessage, session_id: Option>, - mut auth_token: Option, + auth_token: Option, custom_headers: HashMap, max_sse_event_size: usize, ) -> Result< crate::transport::streamable_http_client::StreamableHttpPostResponse, StreamableHttpError, > { - if auth_token.is_none() { - auth_token = Some(self.get_access_token().await?); - } - self.http_client - .post_message_with_max_sse_event_size( - uri, - message, - session_id, - auth_token, - custom_headers, - max_sse_event_size, - ) + 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") + } + + 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 4ed05b4..2ee37aa 100644 --- a/crates/rmcp/src/transport/common/client_side_sse.rs +++ b/crates/rmcp/src/transport/common/client_side_sse.rs @@ -205,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 { @@ -216,6 +221,7 @@ impl Default for ExponentialBackoff { Self { max_times: None, base_duration: Self::DEFAULT_DURATION, + max_delay: None, } } } @@ -227,7 +233,16 @@ impl SseRetryPolicy for ExponentialBackoff { { 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, + }) } } @@ -376,145 +391,157 @@ 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)) => { - 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, - } - } - 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 - .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::Terminated => { - return Poll::Ready(None); - } - }; - // update the state - this.state.set(next_state); - self.poll_next(cx) + 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); + } } } @@ -716,6 +743,34 @@ mod tests { ); } + #[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)); @@ -735,4 +790,75 @@ mod tests { 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, + }; + 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/reqwest/streamable_http_client.rs b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs index e2eeebb..1b47560 100644 --- a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs @@ -94,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) => { @@ -262,6 +293,11 @@ 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}"), ))); @@ -272,17 +308,29 @@ impl StreamableHttpClient for reqwest::Client { 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), + ))), } } _ => { diff --git a/crates/rmcp/src/transport/common/unix_socket.rs b/crates/rmcp/src/transport/common/unix_socket.rs index 5f995db..688c9ad 100644 --- a/crates/rmcp/src/transport/common/unix_socket.rs +++ b/crates/rmcp/src/transport/common/unix_socket.rs @@ -274,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}"), ))); @@ -317,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( diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index 5194ba9..8af045f 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -1,12 +1,16 @@ use std::{ borrow::Cow, - collections::{HashMap, HashSet}, + 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; @@ -19,13 +23,17 @@ use super::common::client_side_sse::{ use crate::{ RoleClient, model::{ - ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, GetMeta, + ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, GetExtensions, GetMeta, InitializedNotification, JsonObject, ProtocolVersion, RequestId, ServerJsonRpcMessage, ServerResult, }, + service::InboundStreamOrigin, transport::{ common::{client_side_sse::SseAutoReconnectStream, mcp_headers}, - worker::{Worker, WorkerQuitReason, WorkerSendRequest, WorkerTransport}, + worker::{ + RequestCancellationRegistration, Worker, WorkerQuitReason, WorkerSendRequest, + WorkerTransport, + }, }, }; @@ -124,7 +132,8 @@ fn negotiate_version_headers( (version, headers) } -#[derive(Debug)] +#[derive(Debug, Error)] +#[error("authorization required: {www_authenticate_header}")] #[non_exhaustive] pub struct AuthRequiredError { pub www_authenticate_header: String, @@ -139,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, @@ -197,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)] @@ -300,12 +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( @@ -456,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 { @@ -475,6 +592,85 @@ impl StreamableHttpClientWorker { } impl StreamableHttpClientWorker { + // 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()), @@ -490,28 +686,16 @@ impl StreamableHttpClientWorker { } } - fn mark_stream_response_pending( - pending_stream_response_ids: &mut HashSet, - request_id: Option, - ) { - if let Some(request_id) = request_id { - pending_stream_response_ids.insert(request_id); - } - } - fn clear_stream_response_pending( pending_stream_response_ids: &mut HashSet, message: &ServerJsonRpcMessage, - ) { - let Some(response_id) = Self::server_response_id(message) else { - return; - }; - if pending_stream_response_ids.remove(response_id) { - return; - } - if let Some(id) = response_id.numeric_string_value() { - pending_stream_response_ids.remove(&RequestId::Number(id)); + ) -> 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( @@ -522,7 +706,8 @@ impl StreamableHttpClientWorker { loop { match sse_worker_rx.try_recv() { Ok(message) => { - Self::clear_stream_response_pending(pending_stream_response_ids, &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(()), @@ -554,6 +739,44 @@ impl StreamableHttpClientWorker { 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 @@ -612,11 +835,36 @@ impl StreamableHttpClientWorker { .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> { @@ -631,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(_) @@ -701,6 +954,7 @@ impl StreamableHttpClientWorker { Self::execute_sse_stream( sse_stream, sse_worker_tx, + InboundStreamOrigin::Unassociated, false, transport_task_ct.child_token(), ) @@ -793,6 +1047,19 @@ impl StreamableHttpClientWorker { 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 } @@ -818,6 +1085,7 @@ impl Worker for StreamableHttpClientWorker { let WorkerSendRequest { responder, message: startup_request, + .. } = context.recv_from_handler().await?; let is_legacy_startup = matches!( &startup_request, @@ -927,17 +1195,31 @@ impl Worker for StreamableHttpClientWorker { clippy::large_enum_variant, reason = "the event is short-lived and boxing would add allocation in the event loop" )] - enum Event { - ClientMessage(WorkerSendRequest), + enum Event { + ClientMessage(WorkerSendRequest>), + ControlMessage(WorkerSendRequest>), + StartPost(WorkerSendRequest>), + PostResult(PostResult), + RecoveryTimeout, ServerMessage(ServerJsonRpcMessage), StreamResult { request_id: Option, - result: Result<(), StreamableHttpError>, + 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 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 { Self::spawn_common_stream( @@ -950,19 +1232,159 @@ impl Worker for StreamableHttpClientWorker { transport_task_ct.clone(), ); } - // Main event loop - capture exit reason so we can do cleanup before returning + // 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?; + 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(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) + }); 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"); @@ -972,48 +1394,90 @@ impl Worker for StreamableHttpClientWorker { }, terminated_stream = streams.join_next(), if !streams.is_empty() => { match terminated_stream { - Some(result) => { - match result { - Ok((request_id, result)) => { - Event::StreamResult { request_id, result } - } - Err(error) => Event::StreamResult { - request_id: None, - result: Err(StreamableHttpError::TokioJoinError(error)), - }, - } - } - 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; - let cancellation_request_id = match &message { - ClientJsonRpcMessage::Notification(notification) => { - match ¬ification.notification { - ClientNotification::CancelledNotification(cancelled) => { - cancelled.params.request_id.clone() - } - _ => None, - } - } - _ => None, - }; - if uses_modern_http && let Some(request_id) = cancellation_request_id { - if let Some(stream_ct) = request_stream_cancellations.remove(&request_id) { - stream_ct.cancel(); - } - pending_stream_response_ids.remove(&request_id); - let _ = responder.send(Ok(())); + 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, + message, ClientJsonRpcMessage::Request(request) if matches!( &request.request, @@ -1022,6 +1486,9 @@ impl Worker for StreamableHttpClientWorker { ); 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. @@ -1084,27 +1551,17 @@ impl Worker for StreamableHttpClientWorker { continue; } - let request_id = Self::client_request_id(&message); - let inline_version = match &message { + 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 is_initialized_notification = matches!( - &message, - ClientJsonRpcMessage::Notification(notification) - if matches!( - ¬ification.notification, - ClientNotification::InitializedNotification(_) - ) - ); - // 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 (request_version, request_headers) = request_version_headers( &protocol_headers, - &message, + message, &negotiated_version, &tool_header_cache, ); @@ -1118,177 +1575,87 @@ impl Worker for StreamableHttpClientWorker { cleanup.protocol_headers = protocol_headers.clone(); } } - let response = self - .client - .post_message_with_max_sse_event_size( - config.uri.clone(), - message.clone(), - session_id.clone(), - config.auth_header.clone(), - request_headers, - config.max_sse_event_size, - ) - .await; - let send_result = match response { - Err(StreamableHttpError::SessionExpired) => { - if let Some(saved_init_request) = saved_init_request - .as_ref() - .filter(|_| config.reinit_on_expired_session) - { - // 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(), - config.max_sse_event_size, - ) - .await - { - Ok(( - new_session_id, - new_negotiated_version, - new_protocol_headers, - )) => { - // Old streams hold the stale session ID. Stop them first - // so no late stale-session messages can arrive after the - // pending requests below are completed. - streams.abort_all(); - while streams.join_next().await.is_some() {} - - // Forward any already queued response messages and fail - // the remaining accepted requests so callers do not wait - // forever for responses that can no longer arrive. - Self::drain_queued_stream_messages( - &mut sse_worker_rx, - &mut context, - &mut pending_stream_response_ids, - ) - .await?; - Self::fail_pending_stream_responses( - &mut context, - &mut pending_stream_response_ids, - ) - .await?; - - session_id = new_session_id; - negotiated_version = new_negotiated_version; - 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 { - Self::spawn_common_stream( - &mut streams, - self.client.clone(), - new_sid.clone(), - &config, - protocol_headers.clone(), - sse_worker_tx.clone(), - transport_task_ct.clone(), - ); - } - - let (_, retry_headers) = request_version_headers( - &protocol_headers, - &message, - &negotiated_version, - &tool_header_cache, - ); - let retry_response = self - .client - .post_message_with_max_sse_event_size( - config.uri.clone(), - message, - session_id.clone(), - config.auth_header.clone(), - retry_headers, - config.max_sse_event_size, - ) - .await; - match retry_response { - Err(e) => Err(e), - Ok(StreamableHttpPostResponse::Accepted) => { - Self::mark_stream_response_pending( - &mut pending_stream_response_ids, - request_id, - ); - tracing::trace!( - "client message accepted after re-init" - ); - Ok(()) - } - Ok(StreamableHttpPostResponse::Json(mut msg, ..)) => { - cache_tools_from_response( - &mut tool_header_cache, - &mut msg, - &negotiated_version, - ); - context.send_to_handler(msg).await?; - Ok(()) - } - Ok(StreamableHttpPostResponse::Sse(stream, ..)) => { - let stream_request_id = request_id.clone(); - Self::mark_stream_response_pending( - &mut pending_stream_response_ids, - 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 stream_ct = transport_task_ct.child_token(); - if uses_modern_http - && let Some(request_id) = - stream_request_id.as_ref() - { - request_stream_cancellations.insert( - request_id.clone(), - stream_ct.clone(), - ); - } - let stream_tx = sse_worker_tx.clone(); - streams.spawn(async move { - let result = Self::execute_sse_stream( - sse_stream, stream_tx, true, stream_ct, - ) - .await; - (stream_request_id, result) - }); - tracing::trace!("got new sse stream after re-init"); - Ok(()) - } - } - } - Err(reinit_err) => Err(reinit_err), - } - } else { - Err(StreamableHttpError::SessionExpired) - } + 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) => { - Self::mark_stream_response_pending( - &mut pending_stream_response_ids, - request_id, - ); tracing::trace!("client message accepted"); Ok(()) } @@ -1296,17 +1663,13 @@ impl Worker for StreamableHttpClientWorker { cache_tools_from_response( &mut tool_header_cache, &mut message, - &negotiated_version, + &version, ); context.send_to_handler(message).await?; Ok(()) } Ok(StreamableHttpPostResponse::Sse(stream, ..)) => { - let stream_request_id = request_id.clone(); - Self::mark_stream_response_pending( - &mut pending_stream_response_ids, - request_id, - ); + let stream_request_id = request_id; let sse_stream = Self::response_sse_to_jsonrpc( stream, session_id.clone(), @@ -1317,16 +1680,38 @@ impl Worker for StreamableHttpClientWorker { config.max_sse_event_size, self.config.retry_config.clone(), ); - let stream_ct = transport_task_ct.child_token(); - if uses_modern_http && let Some(request_id) = stream_request_id.as_ref() + 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(), stream_ct.clone()); + .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::execute_sse_stream( - sse_stream, stream_tx, true, stream_ct, + let result = Self::run_response_stream( + sse_stream, + stream_tx, + origin, + request_ct, + stream_ct, + uses_modern_http, ) .await; (stream_request_id, result) @@ -1355,18 +1740,13 @@ impl Worker for StreamableHttpClientWorker { let _ = responder.send(send_result); } Event::ServerMessage(mut json_rpc_message) => { - if let Some(response_id) = Self::server_response_id(&json_rpc_message) - && let Some(stream_ct) = crate::service::remove_pending_request( - &mut request_stream_cancellations, - response_id, - ) - { - stream_ct.cancel(); - } - Self::clear_stream_response_pending( + // 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, @@ -1385,8 +1765,10 @@ impl Worker for StreamableHttpClientWorker { &mut pending_stream_response_ids, ) .await?; - request_stream_cancellations.remove(&request_id); - if pending_stream_response_ids.remove(&request_id) { + 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( @@ -1407,6 +1789,14 @@ 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 { @@ -1639,6 +2029,19 @@ 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 @@ -1651,15 +2054,17 @@ pub struct StreamableHttpClientTransportConfig { /// [`StreamableHttpClient`] implementations must override the corresponding /// `*_with_max_sse_event_size` methods to enforce it. pub max_sse_event_size: usize, - /// Enables transparent recovery when the server reports an expired session (`HTTP 404`). + /// Automatically creates a new session when the server reports an expired + /// session (`http 404`). /// - /// 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`. + /// 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. /// - /// 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. + /// 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, } @@ -1671,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 @@ -1735,6 +2158,9 @@ 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(), @@ -1751,7 +2177,92 @@ mod tests { use serde_json::json; use super::*; - use crate::model::{ListToolsResult, NumberOrString, ServerResult, Tool}; + 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); @@ -1837,6 +2348,7 @@ mod tests { Arc::new(ExponentialBackoff { max_times: Some(1), base_duration: Duration::ZERO, + max_delay: None, }), ); let mut stream = std::pin::pin!(stream); @@ -1850,6 +2362,133 @@ mod tests { ); } + #[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", @@ -1934,11 +2573,13 @@ mod tests { NumberOrString::String("1".into()), ); - StreamableHttpClientWorker::::clear_stream_response_pending( - &mut pending, - &response, - ); + let matched_id = + StreamableHttpClientWorker::::clear_stream_response_pending( + &mut pending, + &response, + ); + assert_eq!(matched_id, Some(NumberOrString::Number(1))); assert!(pending.is_empty()); } @@ -1949,14 +2590,16 @@ mod tests { let mut pending = HashSet::from([NumberOrString::Number(1), string_id.clone()]); let response = ServerJsonRpcMessage::response( ServerResult::ListToolsResult(ListToolsResult::default()), - string_id, + string_id.clone(), ); - StreamableHttpClientWorker::::clear_stream_response_pending( - &mut pending, - &response, - ); + 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/local.rs b/crates/rmcp/src/transport/streamable_http_server/session/local.rs index 231724b..e03e5b7 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/local.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/local.rs @@ -1032,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")] @@ -1130,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) => { diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index f66235c..1b53820 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/crates/rmcp/src/transport/streamable_http_server/tower.rs @@ -28,10 +28,13 @@ use crate::{ ClientCapabilities, ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorCode, ErrorData, GetExtensions, GetMeta, Implementation, InitializeRequest, InitializeRequestParams, InitializedNotification, JsonObject, JsonRpcError, - ProtocolVersion, RequestId, ServerJsonRpcMessage, + ProtocolVersion, RequestId, ServerConfig, ServerJsonRpcMessage, ServerResult, }, serve_server, - service::{serve_directly_with_ct, uses_legacy_lifecycle}, + service::{ + NotificationContext, RequestContext, Service, negotiate_protocol_version, + serve_directly_with_ct, uses_legacy_lifecycle, + }, transport::{ OneshotTransport, TransportAdapterIdentity, common::{ @@ -52,6 +55,24 @@ use crate::{ 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 { @@ -63,7 +84,7 @@ pub struct StreamableHttpServerConfig { /// When enabled, SSE priming events are sent to enable client reconnection. /// /// Only applies to legacy protocol versions (`< 2026-07-28`). Per SEP-2567, - /// sessions are removed from the `2026-07-28` draft version, so requests + /// 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, @@ -89,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 @@ -124,6 +150,32 @@ pub struct StreamableHttpServerConfig { /// 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 { @@ -142,8 +194,10 @@ impl Default for StreamableHttpServerConfig { 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, } } } @@ -168,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 { @@ -203,12 +263,20 @@ impl StreamableHttpServerConfig { 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: @@ -217,7 +285,7 @@ impl StreamableHttpServerConfig { fn validate_protocol_version_header( headers: &http::HeaderMap, allow_unknown: bool, -) -> Result<(), BoxResponse> { +) -> HttpResult<()> { if let Some(value) = headers.get(HEADER_MCP_PROTOCOL_VERSION) { let version_str = value.to_str().map_err(|_| { Response::builder() @@ -242,7 +310,8 @@ fn validate_protocol_version_header( ))) .boxed(), ) - .expect("valid response")); + .expect("valid response") + .into()); } } Ok(()) @@ -257,48 +326,102 @@ fn message_has_per_request_protocol_version(message: &ClientJsonRpcMessage) -> b } } -#[expect( - clippy::result_large_err, - reason = "BoxResponse is intentionally large; matches other handlers in this file" -)] +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, -) -> Result { +) -> 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 { - if let ClientRequest::InitializeRequest(init) = &req.request { - validate_header_matches_init_body( - headers, - init.params.protocol_version.as_str(), - Some(req.id.clone()), - )?; - } + 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 !matches!(&req.request, ClientRequest::InitializeRequest(_)) - && req - .request - .get_meta() - .missing_required_keys(&ProtocolVersion::V_2026_07_28) - .is_empty() + if req + .request + .get_meta() + .missing_required_keys(&ProtocolVersion::V_2026_07_28) + .is_empty() ); let from_body = match message { - Some(ClientJsonRpcMessage::Request(req)) => match &req.request { - ClientRequest::InitializeRequest(init) => Some(init.params.protocol_version.clone()), - _ => req.request.get_meta().protocol_version(), - }, + Some(ClientJsonRpcMessage::Request(req)) => req.request.get_meta().protocol_version(), _ => None, }; let version = from_body @@ -330,10 +453,10 @@ async fn persist_and_forward_event( output: &mut Option>, ) -> Result<(), EventStoreError> { event.event_id = Some(event_store.store_event(stream_id, &event).await?); - if let Some(sender) = output { - if sender.send(event).await.is_err() { - *output = None; - } + if let Some(sender) = output + && sender.send(event).await.is_err() + { + *output = None; } Ok(()) } @@ -364,16 +487,12 @@ fn invalid_params_jsonrpc_response( .expect("valid response") } -#[expect( - clippy::result_large_err, - reason = "BoxResponse is intentionally large; matches other handlers in this file" -)] /// 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(()); }; @@ -394,19 +513,16 @@ 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(()) } -#[expect( - clippy::result_large_err, - reason = "BoxResponse is intentionally large; matches other handlers in this file" -)] fn validate_request_protocol_version_meta( headers: &HeaderMap, message: &ClientJsonRpcMessage, -) -> Result<(), BoxResponse> { +) -> HttpResult<()> { let ClientJsonRpcMessage::Request(request) = message else { return Ok(()); }; @@ -414,23 +530,33 @@ fn validate_request_protocol_version_meta( return Ok(()); } let is_discover = matches!(&request.request, ClientRequest::DiscoverRequest(_)); - let Some(meta_version) = request.request.get_meta().protocol_version() else { - if is_discover { + 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()), - "Invalid params: server/discover requires protocolVersion in request _meta", - )); + format!( + "Invalid params: request _meta is missing or has malformed required fields: {}", + missing.join(", ") + ), + ) + .into()); } return Ok(()); }; - let Some(header_version) = headers - .get(HEADER_MCP_PROTOCOL_VERSION) - .and_then(|value| value.to_str().ok()) - else { - return Err(invalid_request_jsonrpc_response( + let Some(header_version) = header_version else { + return Err(header_mismatch_jsonrpc_response( Some(request.id.clone()), - "Invalid Request: request _meta protocolVersion requires MCP-Protocol-Version header", - )); + "request _meta protocolVersion requires MCP-Protocol-Version header", + ) + .into()); }; if header_version != meta_version.as_str() { return Err(header_mismatch_jsonrpc_response( @@ -438,11 +564,77 @@ fn validate_request_protocol_version_meta( 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; @@ -452,16 +644,56 @@ fn jsonrpc_http_status(message: &ServerJsonRpcMessage) -> http::StatusCode { match error.error.code { ErrorCode::UNSUPPORTED_PROTOCOL_VERSION | ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY - | ErrorCode::INVALID_PARAMS => http::StatusCode::BAD_REQUEST, + | 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, -) -> Result { +) -> HttpResult { let status = if map_protocol_status { jsonrpc_http_status(&message) } else { @@ -495,15 +727,11 @@ fn header_mismatch_jsonrpc_response( /// 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). -#[expect( - clippy::result_large_err, - reason = "BoxResponse is intentionally large; matches other handlers in this file" -)] fn validate_standard_headers( headers: &HeaderMap, message: &ClientJsonRpcMessage, tool_schema: impl Fn(&str) -> Option>, -) -> Result<(), BoxResponse> { +) -> HttpResult<()> { let version_requires_headers = headers .get(HEADER_MCP_PROTOCOL_VERSION) .and_then(|value| value.to_str().ok()) @@ -536,7 +764,7 @@ fn validate_standard_headers( .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)); + return Err(header_mismatch_jsonrpc_response(request_id, reason).into()); } Ok(()) } @@ -626,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)) @@ -660,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() @@ -694,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 { @@ -722,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(()) } @@ -745,7 +965,7 @@ 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 @@ -833,9 +1053,7 @@ 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. @@ -888,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, } @@ -919,12 +1136,10 @@ 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, @@ -951,19 +1166,18 @@ where tokio::spawn(async move { let mut sender = Some(sender); - if let Some(retry) = retry { - if let Err(error) = persist_and_forward_event( + 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; - } + { + tracing::error!(%stream_id, %error, "failed to persist SSE priming event"); + request_ct.cancel(); + return; } let mut first = first; @@ -1035,7 +1249,7 @@ where service: S, mut request: crate::model::JsonRpcRequest, parts: http::request::Parts, - ) -> Result { + ) -> HttpResult { let peer_info = Self::peer_info_for_stateless_request(&request, &parts.headers); request.request.extensions_mut().insert(parts); let (transport, mut receiver) = @@ -1044,7 +1258,12 @@ where // 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(service, transport, peer_info, request_ct.clone()); + let service = serve_directly_with_ct( + NegotiatingStatelessHttpService(service), + transport, + peer_info, + request_ct.clone(), + ); tokio::spawn(async move { let _ = service.waiting().await; }); @@ -1095,10 +1314,10 @@ where /// 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() { - if let Some(schema) = cache.get(name) { - return schema.clone(); - } + if let Ok(cache) = self.tool_schemas.read() + && let Some(schema) = cache.get(name) + { + return schema.clone(); } let schema = self .get_service() @@ -1288,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( @@ -1329,7 +1540,7 @@ 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 supports_stateless_replay = self.session_manager.event_store().is_some(); @@ -1356,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, @@ -1502,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, @@ -1654,7 +1865,7 @@ where 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( @@ -1672,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 @@ -1743,6 +1954,10 @@ 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) => { @@ -1763,6 +1978,7 @@ where // 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"))?; @@ -1789,8 +2005,12 @@ where // unpersisted response can cancel the in-flight handler on // disconnect (#857). let request_ct = CancellationToken::new(); - let service = - serve_directly_with_ct(service, transport, peer_info, request_ct.clone()); + 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; @@ -1820,7 +2040,8 @@ where std::io::ErrorKind::UnexpectedEof, "no response message received from handler", ), - )); + ) + .into()); }; tracing::trace!(?message); if matches!( @@ -1852,7 +2073,7 @@ 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, @@ -1890,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`, 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_cancelled_response.rs b/crates/rmcp/tests/test_cancelled_response.rs index 5961a7b..864932a 100644 --- a/crates/rmcp/tests/test_cancelled_response.rs +++ b/crates/rmcp/tests/test_cancelled_response.rs @@ -3,13 +3,27 @@ //! 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, - ServerInfo, + ServerConfig, }, service::RequestContext, }; @@ -21,6 +35,7 @@ use tokio::{ 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<()> { @@ -43,7 +58,7 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> { }), ) .await?; - collect_ids_until(&mut reader, 1, READ_TIMEOUT).await?; + collect_ids_until(&mut reader, "1", READ_TIMEOUT).await?; send_json( &mut writer, &json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }), @@ -56,7 +71,7 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> { &mut writer, &json!({ "jsonrpc": "2.0", - "id": 2, + "id": STRING_REQUEST_ID, "method": "tools/call", "params": { "name": "wait-for-cancel", "arguments": {} } }), @@ -67,21 +82,59 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> { &json!({ "jsonrpc": "2.0", "method": "notifications/cancelled", - "params": { "requestId": 2 } + "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 - // an id=2 response is genuine suppression rather than a dead connection. + // 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(&2)); + 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; @@ -91,20 +144,383 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> { struct WaitForCancelServer; impl ServerHandler for WaitForCancelServer { - 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, + 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") { @@ -172,9 +588,9 @@ where /// (then a short grace read to catch any straggler) or the timeout elapses. async fn collect_ids_until( reader: &mut BufReader, - stop_id: u64, + stop_id: &str, timeout: Duration, -) -> anyhow::Result> +) -> anyhow::Result> where R: tokio::io::AsyncRead + Unpin, { @@ -200,8 +616,13 @@ where let Ok(value) = serde_json::from_str::(trimmed) else { continue; }; - if let Some(id) = value.get("id").and_then(Value::as_u64) { - seen.insert(id); + 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); diff --git a/crates/rmcp/tests/test_client_initialization.rs b/crates/rmcp/tests/test_client_initialization.rs index 6c7984c..0ca0182 100644 --- a/crates/rmcp/tests/test_client_initialization.rs +++ b/crates/rmcp/tests/test_client_initialization.rs @@ -51,6 +51,14 @@ async fn client_initialization_accepts_stringified_numeric_response_id() { .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"); } @@ -117,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 index 375364e..5fa2cd1 100644 --- a/crates/rmcp/tests/test_client_lifecycle_modes.rs +++ b/crates/rmcp/tests/test_client_lifecycle_modes.rs @@ -36,11 +36,13 @@ async fn discover_startup_accepts_stringified_numeric_response_id() { }; server .send(ServerJsonRpcMessage::response( - ServerResult::DiscoverResult(DiscoverResult::new( - vec![ProtocolVersion::V_2026_07_28], - ServerCapabilities::default(), - Implementation::new("discover-server", "1.0.0"), - )), + 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 @@ -60,6 +62,110 @@ async fn discover_startup_accepts_stringified_numeric_response_id() { 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); @@ -103,11 +209,13 @@ async fn discover_startup_omits_initialize() { server .send(ServerJsonRpcMessage::response( - ServerResult::DiscoverResult(DiscoverResult::new( - vec![ProtocolVersion::V_2026_07_28], - ServerCapabilities::default(), - Implementation::new("discover-server", "1.0.0"), - )), + 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 @@ -229,6 +337,334 @@ async fn auto_startup_falls_back_after_discover_method_not_found() { 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 = @@ -267,11 +703,13 @@ async fn discover_startup_retries_a_mutually_supported_version() { ); server .send(ServerJsonRpcMessage::response( - ServerResult::DiscoverResult(DiscoverResult::new( - vec![ProtocolVersion::V_2026_07_28], - ServerCapabilities::default(), - Implementation::new("discover-server", "1.0.0"), - )), + 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 @@ -326,11 +764,13 @@ async fn discover_startup_retries_current_version_once_when_server_reports_it_su ); server .send(ServerJsonRpcMessage::response( - ServerResult::DiscoverResult(DiscoverResult::new( - vec![ProtocolVersion::V_2026_07_28], - ServerCapabilities::default(), - Implementation::new("discover-server", "1.0.0"), - )), + 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 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 c3d08cd..8570770 100644 --- a/crates/rmcp/tests/test_deserialization.rs +++ b/crates/rmcp/tests/test_deserialization.rs @@ -70,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!({}))); diff --git a/crates/rmcp/tests/test_discover_http_client_startup.rs b/crates/rmcp/tests/test_discover_http_client_startup.rs index c6051a0..2d92df5 100644 --- a/crates/rmcp/tests/test_discover_http_client_startup.rs +++ b/crates/rmcp/tests/test_discover_http_client_startup.rs @@ -5,11 +5,18 @@ feature = "transport-streamable-http-server" ))] -use std::borrow::Cow; +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::{ClientInfo, DiscoverResult, ErrorCode, ErrorData, ProtocolVersion}, + model::{ClientConfig, DiscoverResult, ErrorCode, ErrorData, ProtocolVersion}, service::{MaybeSendFuture, RequestContext, RoleServer}, transport::{ StreamableHttpClientTransport, @@ -19,6 +26,8 @@ use rmcp::{ }, }, }; +use serde_json::json; +use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; #[derive(Clone, Default)] @@ -46,6 +55,54 @@ impl ServerHandler for LegacyHttpServer { } } +#[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(); @@ -75,7 +132,7 @@ async fn discover_http_client_bootstraps_headers_without_initialize() { let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(format!("http://{address}/mcp")), ); - let client = ClientInfo::default() + let client = ClientConfig::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { @@ -119,7 +176,7 @@ async fn auto_http_client_falls_back_to_stateful_legacy_startup() { let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(format!("http://{address}/mcp")), ); - let client = ClientInfo::default() + let client = ClientConfig::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Auto { @@ -135,3 +192,47 @@ async fn auto_http_client_falls_back_to_stateful_legacy_startup() { 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 b2bc39a..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"), } } @@ -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,7 +1876,7 @@ 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"), } } @@ -1927,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": { @@ -1953,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"), } } 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_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 2b0990d..027a2cc 100644 --- a/crates/rmcp/tests/test_message_schema.rs +++ b/crates/rmcp/tests/test_message_schema.rs @@ -61,15 +61,15 @@ mod tests { ); } - /// The three metadata definitions must expose the MCP 2026-07-28 draft + /// 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`. The keys the draft - /// marks as required stay optional because rmcp generates one schema - /// shared by every supported protocol version; draft-strict validation is + /// 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_draft_schema() { + fn test_metadata_definitions_match_2026_07_28_schema() { let settings = SchemaSettings::draft07(); let schema = settings .into_generator() @@ -80,7 +80,7 @@ mod tests { assert_eq!( definitions["MetaObject"], serde_json::json!({ - "description": "See [specification/draft/basic/index#general-fields] for notes on _meta usage.", + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", "type": "object", "additionalProperties": true, }) @@ -121,6 +121,13 @@ mod tests { 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 fe2eb5e..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 @@ -1206,7 +1206,7 @@ ] }, "MetaObject": { - "description": "See [specification/draft/basic/index#general-fields] for notes on _meta usage.", + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", "type": "object", "additionalProperties": true }, @@ -2233,7 +2233,7 @@ "type": "object", "properties": { "_meta": { - "description": "Protocol-level metadata. Required by the draft wire schema.", + "description": "Protocol-level metadata. Required by the 2026-07-28 wire schema.", "type": "object", "properties": { "io.modelcontextprotocol/clientCapabilities": { @@ -2255,7 +2255,6 @@ "additionalProperties": true, "required": [ "io.modelcontextprotocol/protocolVersion", - "io.modelcontextprotocol/clientInfo", "io.modelcontextprotocol/clientCapabilities" ] }, 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 fe2eb5e..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 @@ -1206,7 +1206,7 @@ ] }, "MetaObject": { - "description": "See [specification/draft/basic/index#general-fields] for notes on _meta usage.", + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", "type": "object", "additionalProperties": true }, @@ -2233,7 +2233,7 @@ "type": "object", "properties": { "_meta": { - "description": "Protocol-level metadata. Required by the draft wire schema.", + "description": "Protocol-level metadata. Required by the 2026-07-28 wire schema.", "type": "object", "properties": { "io.modelcontextprotocol/clientCapabilities": { @@ -2255,7 +2255,6 @@ "additionalProperties": true, "required": [ "io.modelcontextprotocol/protocolVersion", - "io.modelcontextprotocol/clientInfo", "io.modelcontextprotocol/clientCapabilities" ] }, 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 595281e..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 @@ -199,7 +199,7 @@ ] }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -328,7 +328,7 @@ "$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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -759,14 +759,6 @@ } ] }, - "serverInfo": { - "description": "Information about the server implementation.", - "allOf": [ - { - "$ref": "#/definitions/Implementation" - } - ] - }, "supportedVersions": { "description": "Protocol versions implemented by this server.", "type": "array", @@ -785,7 +777,6 @@ "resultType", "supportedVersions", "capabilities", - "serverInfo", "ttlMs", "cacheScope" ] @@ -972,6 +963,10 @@ "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": [ @@ -1144,7 +1139,7 @@ } }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -1789,7 +1784,7 @@ } }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -1850,7 +1845,7 @@ } }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -1911,7 +1906,7 @@ } }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -1971,7 +1966,7 @@ ] }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -2061,7 +2056,7 @@ ] }, "MetaObject": { - "description": "See [specification/draft/basic/index#general-fields] for notes on _meta usage.", + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", "type": "object", "additionalProperties": true }, @@ -2595,7 +2590,7 @@ } }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -3459,6 +3454,14 @@ "SubscriptionsListenResultMeta": { "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": [ + { + "$ref": "#/definitions/Implementation" + } + ] + }, "io.modelcontextprotocol/subscriptionId": { "$ref": "#/definitions/NumberOrString" } 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 595281e..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 @@ -199,7 +199,7 @@ ] }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -328,7 +328,7 @@ "$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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -759,14 +759,6 @@ } ] }, - "serverInfo": { - "description": "Information about the server implementation.", - "allOf": [ - { - "$ref": "#/definitions/Implementation" - } - ] - }, "supportedVersions": { "description": "Protocol versions implemented by this server.", "type": "array", @@ -785,7 +777,6 @@ "resultType", "supportedVersions", "capabilities", - "serverInfo", "ttlMs", "cacheScope" ] @@ -1144,7 +1135,7 @@ } }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -1789,7 +1780,7 @@ } }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -1850,7 +1841,7 @@ } }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -1911,7 +1902,7 @@ } }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -1971,7 +1962,7 @@ ] }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -2061,7 +2052,7 @@ ] }, "MetaObject": { - "description": "See [specification/draft/basic/index#general-fields] for notes on _meta usage.", + "description": "See [MCP general fields](https://modelcontextprotocol.io/specification/2026-07-28/basic#general-fields) for notes on _meta usage.", "type": "object", "additionalProperties": true }, @@ -2595,7 +2586,7 @@ } }, "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/5bed7b30527019e34ccb0eb474636651424501f6/schema/draft/schema.ts#L225-L234", + "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" @@ -3459,6 +3450,14 @@ "SubscriptionsListenResultMeta": { "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": [ + { + "$ref": "#/definitions/Implementation" + } + ] + }, "io.modelcontextprotocol/subscriptionId": { "$ref": "#/definitions/NumberOrString" } diff --git a/crates/rmcp/tests/test_mrtr_behavior.rs b/crates/rmcp/tests/test_mrtr_behavior.rs index cbdd6dd..d9df5af 100644 --- a/crates/rmcp/tests/test_mrtr_behavior.rs +++ b/crates/rmcp/tests/test_mrtr_behavior.rs @@ -13,9 +13,16 @@ use std::sync::{ use rmcp::{ ClientHandler, ServerHandler, + handler::server::{ + tool::{InputResponses as ToolInputResponses, RequestState}, + wrapper::Parameters, + }, model::*, - service::{RequestContext, RoleClient, RoleServer, ServiceError, serve_directly}, + 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: @@ -66,6 +73,54 @@ fn single_elicitation(state: &str) -> InputRequiredResult { 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, @@ -166,8 +221,8 @@ impl MrtrServer { } impl ServerHandler for MrtrServer { - fn get_info(&self) -> ServerInfo { - let mut info = ServerInfo::new( + fn get_info(&self) -> ServerConfig { + let mut info = ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_prompts() @@ -273,28 +328,29 @@ impl ClientHandler for MrtrClient { // Harness // ============================================================================= -fn client_info(protocol_version: ProtocolVersion) -> ClientInfo { - ClientInfo::new( +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) -> ServerInfo { - let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build()); +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: MrtrServer, +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>, { @@ -316,7 +372,7 @@ where let client = serve_directly::( MrtrClient, client_transport, - Some(client_peer_info), + Some(client_peer_info.into()), ); let result = body(client).await; @@ -346,6 +402,23 @@ async fn client_auto_fulfills_input_required_tool_call() -> anyhow::Result<()> { .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(); @@ -521,19 +594,23 @@ async fn request_state_codec_seals_and_verifies_through_the_loop() -> anyhow::Re use rmcp::model::RequestStateCodec; // A shared per-process signing key, mirroring how a real server would derive one. - static KEY: &[u8] = b"integration-signing-key-32-bytes!"; + 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(KEY)) + CODEC.get_or_init(|| RequestStateCodec::new_unchecked(KEY)) } #[derive(Clone, Default)] struct SealingServer; impl ServerHandler for SealingServer { - fn get_info(&self) -> ServerInfo { - let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build()); + fn get_info(&self) -> ServerConfig { + let mut info = ServerConfig::new(ServerCapabilities::builder().enable_tools().build()); info.protocol_version = ProtocolVersion::V_2026_07_28; info } @@ -580,7 +657,7 @@ async fn request_state_codec_seals_and_verifies_through_the_loop() -> anyhow::Re let client = serve_directly::( MrtrClient, client_transport, - Some(server_info(ProtocolVersion::V_2026_07_28)), + Some(server_info(ProtocolVersion::V_2026_07_28).into()), ); let result = client diff --git a/crates/rmcp/tests/test_notification.rs b/crates/rmcp/tests/test_notification.rs index 9aafd3c..dafb0df 100644 --- a/crates/rmcp/tests/test_notification.rs +++ b/crates/rmcp/tests/test_notification.rs @@ -6,7 +6,7 @@ use rmcp::{ ClientHandler, ServerHandler, ServiceExt, model::{ ClientNotification, CustomNotification, ResourceUpdatedNotificationParam, - ServerCapabilities, ServerInfo, ServerNotification, SubscribeRequestParams, + ServerCapabilities, ServerConfig, ServerNotification, SubscribeRequestParams, }, }; use serde_json::json; @@ -16,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_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_resource_not_found_version.rs b/crates/rmcp/tests/test_resource_not_found_version.rs index 44eb363..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, - ReadResourceResponse, + ClientConfig, ErrorCode, ErrorData, InitializeResult, ProtocolVersion, + ReadResourceRequestParams, ReadResourceResponse, }, - service::RequestContext, + service::{RequestContext, serve_directly}, }; #[derive(Debug, Clone, Default)] @@ -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 index 849a661..61240ed 100644 --- a/crates/rmcp/tests/test_result_type_version.rs +++ b/crates/rmcp/tests/test_result_type_version.rs @@ -6,12 +6,12 @@ #![cfg(feature = "client")] use rmcp::{ - ClientHandler, RoleServer, ServerHandler, ServiceExt, + ClientHandler, RoleClient, RoleServer, ServerHandler, model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, ClientInfo, ContentBlock, - ErrorData, ProtocolVersion, ResultType, + CallToolRequestParams, CallToolResponse, CallToolResult, ClientConfig, ContentBlock, + ErrorData, InitializeResult, ProtocolVersion, ResultType, }, - service::RequestContext, + service::{RequestContext, serve_directly}, }; #[derive(Debug, Clone, Default)] @@ -33,27 +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 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 { - ToolServer.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 result = client .call_tool(CallToolRequestParams::new("echo")) diff --git a/crates/rmcp/tests/test_result_type_wire.rs b/crates/rmcp/tests/test_result_type_wire.rs index d5f2015..a22933e 100644 --- a/crates/rmcp/tests/test_result_type_wire.rs +++ b/crates/rmcp/tests/test_result_type_wire.rs @@ -86,6 +86,68 @@ fn legacy_call_tool_result_round_trips_without_result_type() { 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 = diff --git a/crates/rmcp/tests/test_sampling.rs b/crates/rmcp/tests/test_sampling.rs index b108ff4..3c82fd4 100644 --- a/crates/rmcp/tests/test_sampling.rs +++ b/crates/rmcp/tests/test_sampling.rs @@ -375,7 +375,7 @@ fn test_tool_result_content_requires_content() { #[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.14))] +#[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, diff --git a/crates/rmcp/tests/test_sep_2260_request_association.rs b/crates/rmcp/tests/test_sep_2260_request_association.rs index d4e20e1..787d634 100644 --- a/crates/rmcp/tests/test_sep_2260_request_association.rs +++ b/crates/rmcp/tests/test_sep_2260_request_association.rs @@ -1,16 +1,19 @@ #![cfg(all(feature = "server", feature = "client", not(feature = "local")))] -#![allow(deprecated)] +#![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, ClientInfo, ContentBlock, + CallToolRequestParams, CallToolResponse, CallToolResult, ClientConfig, ContentBlock, CreateMessageRequest, CreateMessageRequestParams, CreateMessageResult, ProtocolVersion, - SamplingMessage, ServerCapabilities, ServerInfo, ServerRequest, + SamplingMessage, ServerCapabilities, ServerConfig, ServerRequest, }, - service::RequestContext, + service::{RequestContext, RunningService, serve_directly}, }; use serde_json::{Value, json}; use tokio::{ @@ -18,14 +21,16 @@ use tokio::{ sync::oneshot, }; +type RequestResultSender = oneshot::Sender>; + #[derive(Clone)] struct SamplingServer { - outside: Arc>>>>, + outside: Arc>>, } impl ServerHandler for SamplingServer { - 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( @@ -88,28 +93,51 @@ impl ClientHandler for SamplingClient { .with_stop_reason(CreateMessageResult::STOP_REASON_END_TURN)) } - fn get_info(&self) -> ClientInfo { - let mut info = ClientInfo::default(); + 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 (server_transport, client_transport) = tokio::io::duplex(4096); 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 { - let running = server.serve(server_transport).await?; - running.waiting().await?; + running_server.waiting().await?; anyhow::Ok(()) }); - let client = SamplingClient.serve(client_transport).await?; - let result = client .peer() .call_tool(CallToolRequestParams::new("sample")) @@ -129,19 +157,16 @@ async fn nested_sampling_allowed_standalone_rejected() -> anyhow::Result<()> { #[tokio::test] async fn generic_send_request_bypass_rejected() -> anyhow::Result<()> { - let (server_transport, client_transport) = tokio::io::duplex(4096); 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 { - let running = server.serve(server_transport).await?; - running.waiting().await?; + running_server.waiting().await?; anyhow::Ok(()) }); - let client = SamplingClient.serve(client_transport).await?; - let result = client .peer() .call_tool(CallToolRequestParams::new("sample_generic")) 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 index 824ef0c..e883ccd 100644 --- a/crates/rmcp/tests/test_sep_2260_stream_routing.rs +++ b/crates/rmcp/tests/test_sep_2260_stream_routing.rs @@ -9,7 +9,7 @@ use rmcp::{ ErrorData as McpError, RoleServer, ServerHandler, model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ElicitRequestParams, - ElicitationSchema, ServerCapabilities, ServerInfo, + ElicitationSchema, ServerCapabilities, ServerConfig, }, service::RequestContext, transport::streamable_http_server::{ @@ -23,8 +23,8 @@ use tokio_util::sync::CancellationToken; struct ElicitingServer; impl ServerHandler for ElicitingServer { - 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( diff --git a/crates/rmcp/tests/test_server_discover.rs b/crates/rmcp/tests/test_server_discover.rs index 3f689b9..f2ca3df 100644 --- a/crates/rmcp/tests/test_server_discover.rs +++ b/crates/rmcp/tests/test_server_discover.rs @@ -74,6 +74,153 @@ fn discover_result_deserializes_to_typed_variant() { ); } +#[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( diff --git a/crates/rmcp/tests/test_server_discover_client.rs b/crates/rmcp/tests/test_server_discover_client.rs index adbe577..9046f6b 100644 --- a/crates/rmcp/tests/test_server_discover_client.rs +++ b/crates/rmcp/tests/test_server_discover_client.rs @@ -4,7 +4,7 @@ use rmcp::{ ClientHandler, ServerHandler, ServiceExt, model::{ ClientCapabilities, Implementation, ProtocolVersion, RequestMetaObject, ServerCapabilities, - ServerInfo, + ServerConfig, }, select_protocol_version, }; @@ -13,8 +13,8 @@ use rmcp::{ struct DiscoveryServer; impl ServerHandler for DiscoveryServer { - 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::new("discovery-server", "1.0.0")) } } @@ -70,8 +70,8 @@ async fn client_discover_helper_returns_typed_result() { .expect("discover should succeed"); assert_eq!( - result.server_info, - Implementation::new("discovery-server", "1.0.0") + 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 index 3dbafb1..7762c97 100644 --- a/crates/rmcp/tests/test_server_discover_http.rs +++ b/crates/rmcp/tests/test_server_discover_http.rs @@ -8,7 +8,7 @@ use std::borrow::Cow; use rmcp::{ ServerHandler, - model::{Implementation, ProtocolVersion, ServerCapabilities, ServerInfo}, + model::{Implementation, ProtocolVersion, ServerCapabilities, ServerConfig}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -20,8 +20,8 @@ use tokio_util::sync::CancellationToken; struct DiscoveryServer; impl ServerHandler for DiscoveryServer { - 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::new("discovery-server", "1.0.0")) .with_instructions("Use the tools carefully") } @@ -122,9 +122,11 @@ async fn discover_returns_server_metadata_without_session() { "resultType": "complete", "supportedVersions": ["2025-11-25"], "capabilities": { "tools": {} }, - "serverInfo": { - "name": "discovery-server", - "version": "1.0.0" + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "discovery-server", + "version": "1.0.0" + } }, "instructions": "Use the tools carefully", "ttlMs": 0, @@ -327,6 +329,38 @@ async fn discover_rejects_missing_client_capabilities() { 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; 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 1804b79..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_legacy_session_mode(false) - .with_json_response(true) .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 index 4404e08..9ab302b 100644 --- a/crates/rmcp/tests/test_stateless_server_requests.rs +++ b/crates/rmcp/tests/test_stateless_server_requests.rs @@ -1,14 +1,18 @@ #![cfg(all(feature = "server", not(feature = "local")))] -use std::sync::{Arc, Mutex}; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; use rmcp::{ ServerHandler, ServiceExt, model::{ ClientCapabilities, ClientJsonRpcMessage, ClientRequest, DiscoverRequest, DiscoverRequestParams, ErrorCode, ErrorData, Implementation, ListToolsRequest, - ListToolsResult, PaginatedRequestParams, ProtocolVersion, RequestId, RequestMetaObject, - ServerJsonRpcMessage, + ListToolsResult, NumberOrString, PaginatedRequestParams, ProgressNotificationParam, + ProgressToken, ProtocolVersion, RequestId, RequestMetaObject, ServerJsonRpcMessage, + ServerNotification, }, service::{MaybeSendFuture, RequestContext, RoleServer, ServerInitializeError}, transport::{IntoTransport, Transport}, @@ -162,7 +166,7 @@ async fn stateless_server_uses_each_requests_client_context() { } #[tokio::test] -async fn stateless_server_rejects_malformed_metadata_opener() { +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); @@ -186,6 +190,24 @@ async fn stateless_server_rejects_malformed_metadata_opener() { )) .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"); }; @@ -194,3 +216,73 @@ async fn stateless_server_rejects_malformed_metadata_opener() { 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 2e1e40a..64be3fc 100644 --- a/crates/rmcp/tests/test_stdio_response_concurrency.rs +++ b/crates/rmcp/tests/test_stdio_response_concurrency.rs @@ -6,7 +6,7 @@ use rmcp::{ ErrorData as McpError, ServerHandler, ServiceExt, model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ServerCapabilities, - ServerInfo, + ServerConfig, }, }; use serde_json::{Value, json}; @@ -82,8 +82,8 @@ 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( 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 index 31e7413..05efe21 100644 --- a/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs +++ b/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs @@ -19,7 +19,7 @@ use rmcp::{ ErrorData as McpError, RoleServer, ServerHandler, model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ServerCapabilities, - ServerInfo, + ServerConfig, }, service::RequestContext, transport::streamable_http_server::{ @@ -37,8 +37,8 @@ struct CancelProbe { impl ServerHandler for CancelProbe { #[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( diff --git a/crates/rmcp/tests/test_streamable_http_event_store.rs b/crates/rmcp/tests/test_streamable_http_event_store.rs index cc88f10..979ce7e 100644 --- a/crates/rmcp/tests/test_streamable_http_event_store.rs +++ b/crates/rmcp/tests/test_streamable_http_event_store.rs @@ -19,7 +19,7 @@ use rmcp::{ ErrorData, ServerHandler, model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, - ProgressNotificationParam, ServerCapabilities, ServerInfo, + ProgressNotificationParam, ServerCapabilities, ServerConfig, }, service::RequestContext, transport::streamable_http_server::{ @@ -109,8 +109,8 @@ impl EventStore for InMemoryEventStore { struct ProgressServer; impl ServerHandler for ProgressServer { - 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( 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 b1c09f5..62daa7a 100644 --- a/crates/rmcp/tests/test_streamable_http_json_response.rs +++ b/crates/rmcp/tests/test_streamable_http_json_response.rs @@ -3,7 +3,7 @@ use rmcp::{ ErrorData, ServerHandler, model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, - ProgressNotificationParam, ServerCapabilities, ServerInfo, + ProgressNotificationParam, ServerCapabilities, ServerConfig, }, service::RequestContext, transport::streamable_http_server::{ @@ -57,8 +57,8 @@ const NEGOTIATED_CALL_WITH_PROGRESS_BODY: &str = r#"{ struct ProgressServer; impl ServerHandler for ProgressServer { - 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( 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_protocol_version.rs b/crates/rmcp/tests/test_streamable_http_protocol_version.rs index cf63624..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; @@ -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_stale_session.rs b/crates/rmcp/tests/test_streamable_http_stale_session.rs index be1ac02..c9ea214 100644 --- a/crates/rmcp/tests/test_streamable_http_stale_session.rs +++ b/crates/rmcp/tests/test_streamable_http_stale_session.rs @@ -15,7 +15,7 @@ use http::{HeaderName, HeaderValue}; use rmcp::{ ServiceError, ServiceExt, model::{ - CallToolRequestParams, ClientInfo, ClientJsonRpcMessage, ClientRequest, ErrorCode, + CallToolRequestParams, ClientConfig, ClientJsonRpcMessage, ClientRequest, ErrorCode, ErrorData, InitializeResult, PingRequest, ProtocolVersion, RequestId, ServerCapabilities, ServerJsonRpcMessage, ServerResult, }, @@ -205,7 +205,7 @@ async fn test_reinitialization_completes_accepted_sse_request_instead_of_hanging mock_client, StreamableHttpClientTransportConfig::with_uri("mock://mcp"), ); - let mut client = ClientInfo::default().serve(transport).await?; + let mut client = ClientConfig::default().serve(transport).await?; let peer = client.peer().clone(); let pending_call = tokio::spawn(async move { diff --git a/crates/rmcp/tests/test_streamable_http_standard_headers.rs b/crates/rmcp/tests/test_streamable_http_standard_headers.rs index 510433e..9e70805 100644 --- a/crates/rmcp/tests/test_streamable_http_standard_headers.rs +++ b/crates/rmcp/tests/test_streamable_http_standard_headers.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use rmcp::{ ServerHandler, - model::{ServerCapabilities, ServerInfo, Tool}, + model::{ServerCapabilities, ServerConfig, Tool}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -18,8 +18,8 @@ const SEP_VERSION: &str = "2026-07-28"; struct HeaderValidationServer; impl ServerHandler for HeaderValidationServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) } fn get_tool(&self, name: &str) -> Option { diff --git a/crates/rmcp/tests/test_subscriptions.rs b/crates/rmcp/tests/test_subscriptions.rs index 3f33765..7359530 100644 --- a/crates/rmcp/tests/test_subscriptions.rs +++ b/crates/rmcp/tests/test_subscriptions.rs @@ -19,7 +19,7 @@ use rmcp::{ model::{ ClientNotification, ClientRequest, DiscoverResult, GetMeta, Implementation, NotificationMetaObject, PromptListChangedNotification, ProtocolVersion, ServerCapabilities, - ServerInfo, ServerNotification, ServerResult, SubscriptionFilter, + ServerConfig, ServerNotification, ServerResult, SubscriptionFilter, SubscriptionsAcknowledgedNotification, SubscriptionsAcknowledgedNotificationParams, SubscriptionsListenResult, }, @@ -44,13 +44,14 @@ impl ClientHandler for CountingClient { } impl ServerHandler for ToolsOnlyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + 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( @@ -79,8 +80,8 @@ impl ServerHandler for ToolsOnlyServer { struct ToolsAndPromptsServer; impl ServerHandler for ToolsAndPromptsServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -119,8 +120,8 @@ impl ServerHandler for ToolsAndPromptsServer { struct ResourceSubscriptionServer; impl ServerHandler for ResourceSubscriptionServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_resources() .enable_resources_subscribe() @@ -181,8 +182,8 @@ impl ServerHandler for RemoteCancellationServer { struct FloodServer; impl ServerHandler for FloodServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -216,8 +217,8 @@ struct ClosedSinkServer { struct LeakyServer; impl ServerHandler for LeakyServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -262,8 +263,8 @@ impl rmcp::service::Service for MalformedAcknowledgmentServer { context: RequestContext, ) -> Result { match request { - ClientRequest::DiscoverRequest(_) => { - Ok(ServerResult::DiscoverResult(DiscoverResult::new( + ClientRequest::DiscoverRequest(_) => Ok(ServerResult::DiscoverResult( + DiscoverResult::new( vec![ProtocolVersion::V_2026_07_28], ServerCapabilities::builder() .enable_tools() @@ -271,9 +272,9 @@ impl rmcp::service::Service for MalformedAcknowledgmentServer { .enable_prompts() .enable_prompts_list_changed() .build(), - Implementation::new("malformed-ack-server", "1.0.0"), - ))) - } + ) + .with_server_info(Implementation::new("malformed-ack-server", "1.0.0")), + )), ClientRequest::SubscriptionsListenRequest(_) => { let mut acknowledgment = SubscriptionsAcknowledgedNotification::new( SubscriptionsAcknowledgedNotificationParams::new( @@ -313,14 +314,14 @@ impl rmcp::service::Service for MalformedAcknowledgmentServer { Ok(()) } - fn get_info(&self) -> rmcp::model::ServerInfo { - ServerInfo::default() + fn get_info(&self) -> rmcp::model::ServerConfig { + ServerConfig::default() } } impl ServerHandler for ClosedSinkServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -440,6 +441,13 @@ async fn listen_exposes_acknowledged_filter_and_graceful_result() -> anyhow::Res 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(()) diff --git a/crates/rmcp/tests/test_subscriptions_model.rs b/crates/rmcp/tests/test_subscriptions_model.rs index b0c8e51..8d79fb1 100644 --- a/crates/rmcp/tests/test_subscriptions_model.rs +++ b/crates/rmcp/tests/test_subscriptions_model.rs @@ -1,7 +1,7 @@ use rmcp::model::{ - ClientJsonRpcMessage, ClientRequest, GetMeta, JsonRpcNotification, JsonRpcRequest, - NotificationMetaObject, RequestId, RequestMetaObject, ServerJsonRpcMessage, ServerNotification, - SubscriptionFilter, SubscriptionsAcknowledgedNotification, + ClientJsonRpcMessage, ClientRequest, GetMeta, Implementation, JsonRpcNotification, + JsonRpcRequest, NotificationMetaObject, RequestId, RequestMetaObject, ServerJsonRpcMessage, + ServerNotification, SubscriptionFilter, SubscriptionsAcknowledgedNotification, SubscriptionsAcknowledgedNotificationParams, SubscriptionsListenRequest, SubscriptionsListenRequestParams, SubscriptionsListenResult, SubscriptionsListenResultMeta, }; @@ -155,6 +155,7 @@ fn acknowledged_notification_round_trips_numeric_subscription_id_and_metadata() #[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); @@ -165,6 +166,10 @@ fn listen_result_requires_matching_string_subscription_id_and_preserves_metadata "resultType": "complete", "_meta": { "io.modelcontextprotocol/subscriptionId": "subscription-2", + "io.modelcontextprotocol/serverInfo": { + "name": "test-server", + "version": "1.0.0", + }, "com.example/result": { "reason": "shutdown", }, @@ -178,6 +183,10 @@ fn listen_result_requires_matching_string_subscription_id_and_preserves_metadata 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" })) @@ -215,7 +224,6 @@ fn subscription_schemas_mark_only_draft_required_fields_as_required() { request_schema["properties"]["_meta"]["required"], json!([ "io.modelcontextprotocol/protocolVersion", - "io.modelcontextprotocol/clientInfo", "io.modelcontextprotocol/clientCapabilities" ]) ); @@ -233,5 +241,14 @@ fn subscription_schemas_mark_only_draft_required_fields_as_required() { 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 index 8acfe28..70f102f 100644 --- a/crates/rmcp/tests/test_subscriptions_streamable_http.rs +++ b/crates/rmcp/tests/test_subscriptions_streamable_http.rs @@ -18,8 +18,9 @@ use std::{ use rmcp::{ ClientLifecycleMode, ClientServiceExt, ServerHandler, model::{ - ClientInfo, ClientRequest, ListToolsRequest, ProtocolVersion, RequestMetaObject, - ServerCapabilities, ServerInfo, ServerNotification, SubscriptionFilter, + ClientConfig, ClientRequest, Implementation, ListToolsRequest, ProtocolVersion, + RequestMetaObject, ServerCapabilities, ServerConfig, ServerNotification, + SubscriptionFilter, }, service::{PeerRequestOptions, SubscriptionContext, SubscriptionEnd}, transport::{ @@ -52,13 +53,14 @@ impl ServerHandler for HttpSubscriptionServer { Cow::Borrowed(&[ProtocolVersion::V_2026_07_28, ProtocolVersion::V_2025_11_25]) } - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + 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( @@ -158,7 +160,7 @@ async fn modern_http_listen_uses_post_stream_and_cancels_by_closing_it() -> anyh let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(url.clone()), ); - let client = ClientInfo::default() + let client = ClientConfig::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { @@ -203,7 +205,7 @@ async fn modern_http_graceful_close_returns_final_listen_result() -> anyhow::Res let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(url), ); - let client = ClientInfo::default() + let client = ClientConfig::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { @@ -217,10 +219,16 @@ async fn modern_http_graceful_close_returns_final_listen_result() -> anyhow::Res assert!(subscription.next().await?.is_some()); assert!(subscription.next().await?.is_none()); - assert!(matches!( - subscription.end(), - Some(SubscriptionEnd::Graceful(_)) - )); + 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(); @@ -233,7 +241,7 @@ async fn modern_http_stream_close_without_result_is_abrupt() -> anyhow::Result<( let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(url), ); - let client = ClientInfo::default() + let client = ClientConfig::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { @@ -266,7 +274,7 @@ async fn modern_http_lifecycle_stays_sessionless_for_older_application_version() let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(url), ); - let client = ClientInfo::default() + let client = ClientConfig::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { diff --git a/crates/rmcp/tests/test_task.rs b/crates/rmcp/tests/test_task.rs index ea1a259..0178bab 100644 --- a/crates/rmcp/tests/test_task.rs +++ b/crates/rmcp/tests/test_task.rs @@ -13,9 +13,9 @@ use rmcp::{ use serde_json::json; #[derive(Debug, serde::Deserialize, rmcp::schemars::JsonSchema)] -pub struct SumArgs { - pub a: i32, - pub b: i32, +struct SumArgs { + a: i32, + b: i32, } #[derive(Clone)] @@ -107,8 +107,8 @@ impl ServerHandler for TaskServer { self.tasks.cancel_task(&request.task_id) } - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_tasks() @@ -117,8 +117,8 @@ impl ServerHandler for TaskServer { } } -fn tasks_client_info() -> ClientInfo { - ClientInfo::new( +fn tasks_client_info() -> ClientConfig { + ClientConfig::new( ClientCapabilities::builder().enable_tasks().build(), Implementation::from_build_env(), ) @@ -275,8 +275,8 @@ impl ServerHandler for AlwaysTaskServer { Ok(CallToolResponse::Task(CreateTaskResult::new(task))) } - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_tasks() diff --git a/crates/rmcp/tests/test_tool_disable_notification.rs b/crates/rmcp/tests/test_tool_disable_notification.rs index cd87805..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::{CallToolResponse, CallToolResult, ServerCapabilities, ServerInfo, Tool}, + model::{CallToolResponse, CallToolResult, ServerCapabilities, ServerConfig, Tool}, service::{MaybeSendFuture, NotificationContext}, }; use tokio::sync::{Notify, RwLock}; @@ -41,8 +41,8 @@ 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( diff --git a/crates/rmcp/tests/test_tool_macros.rs b/crates/rmcp/tests/test_tool_macros.rs index 9b9530a..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() } } @@ -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_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_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/docs/OAUTH_SUPPORT.md b/docs/OAUTH_SUPPORT.md index 1d09d74..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 Authorization Specification](https://modelcontextprotocol.io/specification/draft/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; @@ -158,7 +161,7 @@ distinguish server-published metadata from synthesized metadata. 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/draft/basic/authorization/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 @@ -250,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 @@ -272,6 +295,130 @@ 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 - **Authorization Code client**: [`examples/clients/src/auth/oauth_client.rs`](../examples/clients/src/auth/oauth_client.rs) @@ -329,7 +476,7 @@ If you encounter authorization issues, check the following: ## References -- [MCP Authorization Specification](https://modelcontextprotocol.io/specification/draft/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/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 58565fb..8fa18df 100644 --- a/examples/clients/src/auth/oauth_client.rs +++ b/examples/clients/src/auth/oauth_client.rs @@ -8,8 +8,9 @@ use axum::{ routing::get, }; use rmcp::{ - ServiceExt, - model::ClientInfo, + RoleClient, ServiceExt, + model::ClientConfig, + service::RunningService, transport::{ StreamableHttpClientTransport, auth::{AuthClient, AuthorizationRequest, OAuthState}, @@ -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,73 +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 no scopes lets the SDK auto-select from the server's - // WWW-Authenticate header, 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), - ) - .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 index 7ab612b..9c04531 100644 --- a/examples/clients/src/subscriptions_streamhttp.rs +++ b/examples/clients/src/subscriptions_streamhttp.rs @@ -1,6 +1,6 @@ use rmcp::{ ClientLifecycleMode, ClientServiceExt, - model::{ClientInfo, ProtocolVersion, SubscriptionFilter}, + model::{ClientConfig, ProtocolVersion, SubscriptionFilter}, transport::StreamableHttpClientTransport, }; @@ -13,7 +13,7 @@ async fn main() -> anyhow::Result<()> { .init(); let transport = StreamableHttpClientTransport::from_uri("http://127.0.0.1:8000/mcp"); - let client = ClientInfo::default() + let client = ClientConfig::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { diff --git a/examples/clients/src/task_stdio.rs b/examples/clients/src/task_stdio.rs index 465cb70..a946f75 100644 --- a/examples/clients/src/task_stdio.rs +++ b/examples/clients/src/task_stdio.rs @@ -33,7 +33,7 @@ async fn main() -> Result<()> { .init(); // Declare the tasks extension in our client capabilities (SEP-2663). - let client_info = rmcp::model::ClientInfo::new( + let client_info = rmcp::model::ClientConfig::new( ClientCapabilities::builder().enable_tasks().build(), rmcp::model::Implementation::from_build_env(), ); 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 ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_prompts() .enable_resources() @@ -270,7 +270,7 @@ impl ServerHandler for Counter { async fn initialize( &self, - _request: InitializeRequestParams, + request: InitializeRequestParams, context: RequestContext, ) -> Result { if let Some(http_request_part) = context.extensions.get::() { @@ -278,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) } } 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 4d8d53a..30aed17 100644 --- a/examples/servers/src/common/task_demo.rs +++ b/examples/servers/src/common/task_demo.rs @@ -151,8 +151,8 @@ impl ServerHandler for TaskDemo { self.tasks.cancel_task(&request.task_id) } - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_tasks() 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 d506a9c..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( @@ -145,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 index 40b53c3..b6c87b6 100644 --- a/examples/servers/src/mrtr.rs +++ b/examples/servers/src/mrtr.rs @@ -38,30 +38,83 @@ use rmcp::{ model::*, service::{RequestContext, RoleClient, RoleServer}, }; +use serde::{Deserialize, Serialize}; use serde_json::json; -/// A stable, high-entropy secret. In a real deployment, load this from your -/// secret manager and keep it out of clients' reach. It must stay constant for +/// 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 Default for WeatherServer { - fn default() -> Self { - Self { - codec: RequestStateCodec::new(REQUEST_STATE_KEY), - } +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) -> ServerInfo { - let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build()); + 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 @@ -78,7 +131,9 @@ impl ServerHandler for WeatherServer { None => { let sealed = self .codec - .seal_json(&json!({ "awaiting": "city" })) + .seal_json(&WeatherRequestState { + awaiting: WeatherInput::City, + }) .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; let mut input_requests = InputRequests::new(); @@ -103,21 +158,11 @@ impl ServerHandler for WeatherServer { // Retry round: verify the echoed state before trusting it, read the // elicited city, and return the final result. Some(sealed) => { - let _state: serde_json::Value = self.codec.open_json(&sealed).map_err(|_| { - ErrorData::invalid_params("tampered or unknown request state", None) - })?; - - let city = request - .input_responses - .as_ref() - .and_then(|r| r.get("city")) - .and_then(|v| v["content"]["city"].as_str()) - .unwrap_or("your area"); - - Ok(CallToolResult::success(vec![ContentBlock::text(format!( - "It is sunny in {city}." - ))]) - .into()) + let state = self.open_request_state(&sealed)?; + match state.awaiting { + WeatherInput::City => finish_weather_request(request.input_responses.as_ref()) + .map(CallToolResponse::from), + } } } } @@ -128,8 +173,8 @@ impl ServerHandler for WeatherServer { struct InteractiveClient; impl ClientHandler for InteractiveClient { - fn get_info(&self) -> ClientInfo { - ClientInfo::new( + fn get_info(&self) -> ClientConfig { + ClientConfig::new( ClientCapabilities::builder().enable_elicitation().build(), Implementation::new("mrtr-example-client", env!("CARGO_PKG_VERSION")), ) @@ -152,10 +197,11 @@ impl ClientHandler for InteractiveClient { #[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 = WeatherServer::default() + let server = server .serve(server_transport) .await .expect("server should start"); @@ -163,7 +209,7 @@ async fn main() -> anyhow::Result<()> { }); // Connect the client (this performs the initialize handshake). - let client = InteractiveClient::default().serve(client_transport).await?; + 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) =="); @@ -197,3 +243,91 @@ async fn main() -> anyhow::Result<()> { 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 be230ad..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. ", diff --git a/examples/servers/src/subscriptions_streamhttp.rs b/examples/servers/src/subscriptions_streamhttp.rs index aead85f..7615d59 100644 --- a/examples/servers/src/subscriptions_streamhttp.rs +++ b/examples/servers/src/subscriptions_streamhttp.rs @@ -2,7 +2,7 @@ use std::{borrow::Cow, time::Duration}; use rmcp::{ ErrorData, ServerHandler, - model::{ProtocolVersion, ServerCapabilities, ServerInfo, SubscriptionFilter}, + model::{ProtocolVersion, ServerCapabilities, ServerConfig, SubscriptionFilter}, service::SubscriptionContext, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, @@ -18,8 +18,8 @@ impl ServerHandler for SubscriptionServer { Cow::Borrowed(&[ProtocolVersion::V_2026_07_28]) } - fn get_info(&self) -> ServerInfo { - ServerInfo::new( + fn get_info(&self) -> ServerConfig { + ServerConfig::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() 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"