diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 168b680..bdd4faf 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -16,15 +16,11 @@ updates: - package-ecosystem: "github-actions" directory: "/" schedule: - interval: "weekly" + interval: "daily" 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" + include: "scope" \ No newline at end of file diff --git a/.github/workflows/auto-label-pr.yml b/.github/workflows/auto-label-pr.yml index d91ce04..9850daf 100644 --- a/.github/workflows/auto-label-pr.yml +++ b/.github/workflows/auto-label-pr.yml @@ -3,9 +3,6 @@ on: # Runs workflow when activity on a PR in the workflow's repository occurs. pull_request_target: -permissions: - contents: read - jobs: auto-label: permissions: @@ -23,7 +20,7 @@ jobs: PR_URL: ${{ github.event.pull_request.html_url }} steps: - - uses: actions/labeler@bf12e9b00b37c5c0ca2b87b79b2daf7891dbda13 # v7 + - uses: actions/labeler@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 7ffafe1..ad6cad1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -12,35 +12,34 @@ 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@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 with: fetch-depth: 0 - name: Setup Node.js - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7 + uses: actions/setup-node@v7 with: node-version: '22' - name: Install commitlint - run: npm install --global @commitlint/cli@20.4.3 @commitlint/config-conventional@20.4.3 - + run: | + npm install --save-dev @commitlint/cli @commitlint/config-conventional + echo "module.exports = {extends: ['@commitlint/config-conventional']}" > commitlint.config.js + - name: Lint commit messages - run: commitlint --extends @commitlint/config-conventional --from ${{ github.event.pull_request.base.sha }} --to ${{ github.event.pull_request.head.sha }} --verbose + run: npx commitlint --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@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 - name: Install Rust fmt run: rustup toolchain install nightly --component rustfmt @@ -52,40 +51,32 @@ jobs: name: Lint with Clippy runs-on: ubuntu-latest steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 - name: Install Rust - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@stable - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - - name: Run clippy (all features) + - name: Run clippy 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@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 with: fetch-depth: 0 - name: Install Rust - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@stable - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Install cargo-semver-checks - uses: taiki-e/install-action@3f74d7c16a4242f1c95561e98edc25d36adb4375 # v2.87.12 + uses: taiki-e/install-action@v2.85.11 with: tool: cargo-semver-checks @@ -126,7 +117,7 @@ jobs: runs-on: ubuntu-latest if: github.event_name == 'pull_request' steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 with: fetch-depth: 0 @@ -134,12 +125,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@7c8d7d138f5c09cef361f8214cf96882cd029cdb # nightly + uses: dtolnay/rust-toolchain@nightly - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Install cargo-public-api - uses: taiki-e/install-action@3f74d7c16a4242f1c95561e98edc25d36adb4375 # v2.87.12 + uses: taiki-e/install-action@v2.85.11 with: tool: cargo-public-api @@ -210,9 +201,9 @@ jobs: name: spell check with typos runs-on: ubuntu-latest steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 - name: Spell Check Repo - uses: crate-ci/typos@d43b6c087ac471e2ea7b8af622ff15f05c0c365b # master + uses: crate-ci/typos@master msrv: name: Check MSRV @@ -220,7 +211,7 @@ jobs: permissions: contents: read steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 - name: Read MSRV from Cargo.toml run: | @@ -229,11 +220,11 @@ jobs: echo "MSRV=$MSRV" >> "$GITHUB_ENV" - name: Install MSRV Rust - uses: dtolnay/rust-toolchain@6c977a6ca4077a0ceb28ffbe03f59d46e9ac8772 # master + uses: dtolnay/rust-toolchain@master with: toolchain: ${{ env.MSRV }} - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Check default workspace members with MSRV run: cargo +${{ env.MSRV }} check --all-targets --all-features @@ -242,19 +233,19 @@ jobs: name: Run Tests runs-on: ubuntu-latest steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 # install nodejs - name: Setup Node.js - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7 + uses: actions/setup-node@v7 with: node-version: '22' - name: Install uv - uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 + uses: astral-sh/setup-uv@v7 - name: Install Rust - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@stable - name: Set up Python run: uv python install @@ -262,7 +253,7 @@ jobs: - name: Create venv for python run: uv venv - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Run tests run: cargo test --all-features @@ -271,19 +262,19 @@ jobs: name: Run Tests (no local feature) runs-on: ubuntu-latest steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 # install nodejs - name: Setup Node.js - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7 + uses: actions/setup-node@v7 with: node-version: '22' - name: Install uv - uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 + uses: astral-sh/setup-uv@v7 - name: Install Rust - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@stable - name: Set up Python run: uv python install @@ -291,7 +282,7 @@ jobs: - name: Create venv for python run: uv venv - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Run tests without local feature run: | @@ -305,21 +296,21 @@ jobs: name: Code Coverage runs-on: ubuntu-latest permissions: - contents: read + contents: write steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 # install nodejs - name: Setup Node.js - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7 + uses: actions/setup-node@v7 with: node-version: '22' - name: Install uv - uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 + uses: astral-sh/setup-uv@v7 - name: Install Rust - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@stable - name: Set up Python run: uv python install @@ -327,7 +318,7 @@ jobs: - name: Create venv for python run: uv venv - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Install cargo-llvm-cov run: cargo install cargo-llvm-cov @@ -342,19 +333,19 @@ jobs: name: Example test runs-on: ubuntu-latest steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 # install nodejs - name: Setup Node.js - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7 + uses: actions/setup-node@v7 with: node-version: '22' - name: Install uv - uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 + uses: astral-sh/setup-uv@v7 - name: Install Rust - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@stable - name: Set up Python run: uv python install @@ -362,7 +353,7 @@ jobs: - name: Create venv for python run: uv venv - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Add target WASI preview 2 run: | @@ -398,12 +389,12 @@ jobs: name: Security Audit runs-on: ubuntu-latest steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 - name: Install Rust - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@stable - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Install cargo-audit run: cargo install cargo-audit @@ -415,12 +406,12 @@ jobs: name: Generate Documentation runs-on: ubuntu-latest steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 - name: Install Rust - uses: dtolnay/rust-toolchain@7c8d7d138f5c09cef361f8214cf96882cd029cdb # nightly + uses: dtolnay/rust-toolchain@nightly - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Generate documentation run: | @@ -441,7 +432,7 @@ jobs: # This happened recently in the attack on `tj-actions/changed-files`, but # has happened many times before as well. - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 - name: Update Rust run: | diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index aa77b0a..e8be9ff 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -8,9 +8,6 @@ on: schedule: - cron: '0 0 * * 1' # Weekly on Monday -permissions: - contents: read - jobs: analyze: name: Analyze @@ -24,16 +21,16 @@ jobs: language: [rust, javascript-typescript, python, actions] steps: - name: Checkout repository - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + uses: actions/checkout@v7 - name: Initialize CodeQL - uses: github/codeql-action/init@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 + uses: github/codeql-action/init@v4.37.6 with: languages: ${{ matrix.language }} config-file: ./.github/codeql/codeql-config.yml - name: Autobuild - uses: github/codeql-action/autobuild@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 + uses: github/codeql-action/autobuild@v4.37.6 - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 + uses: github/codeql-action/analyze@v4.37.6 diff --git a/.github/workflows/conformance.yml b/.github/workflows/conformance.yml index 1233deb..ced1875 100644 --- a/.github/workflows/conformance.yml +++ b/.github/workflows/conformance.yml @@ -14,9 +14,6 @@ concurrency: env: CONFORMANCE_VERSION: "0.2.0-alpha.10" -permissions: - contents: read - jobs: server: runs-on: ubuntu-latest @@ -24,12 +21,12 @@ jobs: permissions: contents: read steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@stable - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Build conformance binaries run: cargo build -p mcp-conformance @@ -94,7 +91,7 @@ jobs: - name: Upload results if: always() - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + uses: actions/upload-artifact@v7 with: name: conformance-server-results path: | @@ -107,12 +104,12 @@ jobs: permissions: contents: read steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@stable - - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2 + - uses: Swatinem/rust-cache@v2 - name: Build conformance binaries run: cargo build -p mcp-conformance @@ -143,7 +140,7 @@ jobs: - name: Upload results if: always() - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + uses: actions/upload-artifact@v7 with: name: conformance-client-results path: conformance-client-results diff --git a/.github/workflows/release-plz.yml b/.github/workflows/release-plz.yml index 311b74e..04f5119 100644 --- a/.github/workflows/release-plz.yml +++ b/.github/workflows/release-plz.yml @@ -1,6 +1,8 @@ name: Release-plz -permissions: {} +permissions: + pull-requests: write + contents: write on: push: @@ -16,20 +18,19 @@ jobs: if: ${{ github.repository_owner == 'modelcontextprotocol' }} permissions: contents: write - pull-requests: read steps: - name: Checkout repository - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + uses: actions/checkout@v7 with: fetch-depth: 0 - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@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 --rev c2ea743bacf9bbd56e344df68219e60b0ad35d04 release-plz + run: cargo install --locked --git https://github.com/DaleSeo/release-plz --branch feat/semver-check-features release-plz - name: Run release-plz release run: release-plz release env: @@ -50,17 +51,17 @@ jobs: cancel-in-progress: false steps: - name: Checkout repository - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + uses: actions/checkout@v7 with: fetch-depth: 0 - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + uses: dtolnay/rust-toolchain@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 --rev c2ea743bacf9bbd56e344df68219e60b0ad35d04 release-plz + run: cargo install --locked --git https://github.com/DaleSeo/release-plz --branch feat/semver-check-features 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 deleted file mode 100644 index 453c0c5..0000000 --- a/.github/workflows/scorecard.yml +++ /dev/null @@ -1,48 +0,0 @@ -# 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 fe1baa9..58794d5 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -4,9 +4,6 @@ on: - cron: "0 0 * * *" workflow_dispatch: -permissions: - contents: read - jobs: stale: name: Clean Up @@ -15,7 +12,7 @@ jobs: issues: write pull-requests: write steps: - - uses: actions/stale@4391f3da665fdf50b6810c1a66712fb9ba21aa93 # v11 + - uses: actions/stale@v11 with: days-before-stale: 60 days-before-close: 14 diff --git a/.github/workflows/triage.yml b/.github/workflows/triage.yml index 4f1b9ab..d0d9186 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@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/checkout@v7 - name: Install jq run: sudo apt-get install -y jq diff --git a/Cargo.toml b/Cargo.toml index 0629e82..946255d 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.4.0", path = "./crates/rmcp" } -rmcp-macros = { version = "3.4.0", path = "./crates/rmcp-macros" } +rmcp = { version = "3.1.3", path = "./crates/rmcp" } +rmcp-macros = { version = "3.1.3", path = "./crates/rmcp-macros" } [workspace.package] edition = "2024" rust-version = "1.88" -version = "3.4.0" +version = "3.1.3" authors = ["4t145 "] license = "Apache-2.0" repository = "https://github.com/modelcontextprotocol/rust-sdk/" diff --git a/README.md b/README.md index 7b24d6a..6ab22d6 100644 --- a/README.md +++ b/README.md @@ -99,10 +99,10 @@ Use [`ClientServiceExt::serve_with_lifecycle`](crates/rmcp/src/service/client.rs select another lifecycle explicitly: ```rust, ignore -use rmcp::{ClientLifecycleMode, ClientServiceExt, ProtocolVersion, model::ClientConfig}; +use rmcp::{ClientInfo, ClientLifecycleMode, ClientServiceExt, ProtocolVersion}; // Start directly with server/discover and include client metadata on every request. -let client = ClientConfig::default() +let client = ClientInfo::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { @@ -113,7 +113,7 @@ let client = ClientConfig::default() // Or probe the discover lifecycle and fall back when a legacy server reports // that server/discover is not implemented or does not respond within 10 seconds. -let client = ClientConfig::default() +let client = ClientInfo::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Auto { @@ -369,8 +369,8 @@ use serde_json::json; struct MyServer; impl ServerHandler for MyServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_resources() .build(), @@ -582,8 +582,8 @@ impl MyServer { #[prompt_handler] impl ServerHandler for MyServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_prompts().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_prompts().build()) } } ``` @@ -934,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) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_logging() .build(), @@ -1007,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) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_completions() .enable_prompts() @@ -1253,8 +1253,8 @@ use rmcp::{ }; impl ServerHandler for MyServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -1363,21 +1363,11 @@ async fn call_tool(&self, request: CallToolRequestParams, _ctx: RequestContext **`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. +> **`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. ### Client-side @@ -1584,7 +1574,7 @@ use rmcp::transport::StreamableHttpClientTransport; // Defaults are stateless-friendly. let transport = StreamableHttpClientTransport::from_uri("http://localhost:8000/mcp"); -let client = ClientConfig::default().serve(transport).await?; +let client = ClientInfo::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) @@ -1638,27 +1628,9 @@ server example). The client transport connects with a single URI: use rmcp::transport::StreamableHttpClientTransport; let transport = StreamableHttpClientTransport::from_uri("http://localhost:8000/mcp"); -let client = ClientConfig::default().serve(transport).await?; +let client = ClientInfo::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 @@ -1742,10 +1714,10 @@ knows what the other supports. Declare yours with the `ServerCapabilities` builder in `get_info()`: ```rust,ignore -use rmcp::model::{ServerCapabilities, ServerConfig}; +use rmcp::model::{ServerCapabilities, ServerInfo}; -fn get_info(&self) -> ServerConfig { - ServerConfig::new( +fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_prompts() diff --git a/conformance/Cargo.toml b/conformance/Cargo.toml index e8de6f0..fc4be6d 100644 --- a/conformance/Cargo.toml +++ b/conformance/Cargo.toml @@ -19,7 +19,6 @@ 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", @@ -32,7 +31,6 @@ 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/src/bin/client.rs b/conformance/src/bin/client.rs index aef2c27..5d65410 100644 --- a/conformance/src/bin/client.rs +++ b/conformance/src/bin/client.rs @@ -1,10 +1,3 @@ -#![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::*, @@ -13,8 +6,7 @@ use rmcp::{ AuthClient, AuthorizationManager, StreamableHttpClientTransport, auth::{ AuthorizationCallback, AuthorizationRequest, ClientCredentialsConfig, - InMemoryCredentialStore, JwtSigningAlgorithm, OAuthState, default_oauth_http_client, - enterprise::{EmaAuthorizationServer, EmaClientAuthentication, EmaExchangeRequest}, + InMemoryCredentialStore, JwtSigningAlgorithm, OAuthState, }, streamable_http_client::StreamableHttpClientTransportConfig, }, @@ -44,17 +36,6 @@ 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 { @@ -74,8 +55,8 @@ impl ClientHandler for BasicClientHandler {} struct ElicitationDefaultsClientHandler; impl ClientHandler for ElicitationDefaultsClientHandler { - fn get_info(&self) -> ClientConfig { - let mut info = ClientConfig::default(); + fn get_info(&self) -> ClientInfo { + let mut info = ClientInfo::default(); info.capabilities.elicitation = Some( ElicitationCapability::new() .with_form(FormElicitationCapability::new().with_schema_validation(true)), @@ -163,8 +144,8 @@ impl ClientHandler for ElicitationDefaultsClientHandler { struct FullClientHandler; impl ClientHandler for FullClientHandler { - fn get_info(&self) -> ClientConfig { - let mut info = ClientConfig::default(); + fn get_info(&self) -> ClientInfo { + let mut info = ClientInfo::default(); info.capabilities.elicitation = Some( ElicitationCapability::new() .with_form(FormElicitationCapability::new().with_schema_validation(true)), @@ -779,66 +760,6 @@ 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, @@ -1189,11 +1110,6 @@ 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 5c4b5c8..3b552d7 100644 --- a/conformance/src/bin/server.rs +++ b/conformance/src/bin/server.rs @@ -52,10 +52,6 @@ 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 { @@ -74,7 +70,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_unchecked(REQUEST_STATE_KEY), + request_state_codec: RequestStateCodec::new(REQUEST_STATE_KEY), tasks: TaskManager::new(), } } @@ -704,8 +700,8 @@ impl ServerHandler for ConformanceServer { (name == "test_custom_header").then(custom_header_tool) } - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_prompts() .enable_prompts_list_changed() @@ -1744,9 +1740,7 @@ async fn main() -> anyhow::Result<()> { tracing::info!("Starting conformance server on {}", bind_addr); let server = ConformanceServer::new(); - let config = StreamableHttpServerConfig::default() - .with_allowed_origins([format!("http://{bind_addr}")]) - .enforce_origin_validation(); + let config = StreamableHttpServerConfig::default(); 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 6e35cf3..910257a 100644 --- a/crates/rmcp-macros/CHANGELOG.md +++ b/crates/rmcp-macros/CHANGELOG.md @@ -7,32 +7,6 @@ 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 diff --git a/crates/rmcp-macros/src/lib.rs b/crates/rmcp-macros/src/lib.rs index 30433a8..156e53b 100644 --- a/crates/rmcp-macros/src/lib.rs +++ b/crates/rmcp-macros/src/lib.rs @@ -56,7 +56,6 @@ 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 /// @@ -123,37 +122,6 @@ 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()) @@ -217,8 +185,8 @@ pub fn tool_router(attr: TokenStream, input: TokenStream) -> TokenStream { /// ```rust,ignore /// #[tool_handler] /// impl ServerHandler for MyToolHandler { -/// fn get_info(&self) -> ServerConfig { -/// ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) +/// fn get_info(&self) -> ServerInfo { +/// ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) /// } /// } /// ``` diff --git a/crates/rmcp-macros/src/tool_handler.rs b/crates/rmcp-macros/src/tool_handler.rs index c9f4a25..e274c10 100644 --- a/crates/rmcp-macros/src/tool_handler.rs +++ b/crates/rmcp-macros/src/tool_handler.rs @@ -116,7 +116,7 @@ pub(crate) enum CallerCapability { Prompts, } -/// Build a `get_info()` method that returns `ServerConfig` with the appropriate capabilities. +/// Build a `get_info()` method that returns `ServerInfo` with the appropriate capabilities. /// /// The caller declares its own capability via `caller`. Sibling handler attributes /// (`prompt_handler`, `tool_handler`) are detected automatically @@ -157,8 +157,8 @@ pub(crate) fn build_get_info( } syn::parse2::(quote! { - fn get_info(&self) -> rmcp::model::ServerConfig { - rmcp::model::ServerConfig::new( + fn get_info(&self) -> rmcp::model::ServerInfo { + rmcp::model::ServerInfo::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 6efd692..edc8630 100644 --- a/crates/rmcp-macros/src/tool_router.rs +++ b/crates/rmcp-macros/src/tool_router.rs @@ -17,8 +17,6 @@ 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 { @@ -27,7 +25,6 @@ impl Default for ToolRouterAttribute { router: format_ident!("tool_router"), vis: None, server_handler: false, - allow_empty: false, } } } @@ -38,7 +35,6 @@ pub fn tool_router(attr: TokenStream, input: TokenStream) -> syn::Result(input)?; // find all function marked with `#[rmcp::tool]` @@ -62,19 +58,6 @@ 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 a3e4ab7..366ed5d 100644 --- a/crates/rmcp/CHANGELOG.md +++ b/crates/rmcp/CHANGELOG.md @@ -7,63 +7,6 @@ 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 diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index ac698f9..9dd27a2 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -18,7 +18,6 @@ exhaustive_enums = "warn" features = [ "auth", "auth-client-credentials-jwt", - "auth-enterprise-managed", "base64", "client", "client-side-sse", @@ -71,7 +70,6 @@ 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 = [ @@ -88,7 +86,7 @@ url = { version = "2.4", optional = true } tower-service = { version = "0.3", optional = true } # for child process transport -process-wrap = { version = "10.0", features = ["tokio1"], optional = true } +process-wrap = { version = "9.0", features = ["tokio1"], optional = true } # for cross-platform executable path resolution which = { version = "8", optional = true } @@ -130,7 +128,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", "dep:zeroize", "base64"] +request-state = ["dep:hmac", "dep:sha2", "base64"] # reqwest http client __reqwest = ["dep:reqwest"] @@ -197,11 +195,10 @@ transport-streamable-http-server-session = [ tower = ["dep:tower-service"] 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", "test-util"] } +tokio = { version = "1", features = ["full"] } schemars = { version = "1.1.0", features = ["chrono04"] } axum = { version = "0.8", default-features = false, features = ["http1", "tokio"] } hyper = { version = "1", features = ["server", "http1"] } @@ -215,7 +212,7 @@ tracing-subscriber = { version = "0.3", features = [ "fmt", ] } async-trait = "0.1" -rstest = "0.27.0" +rstest = "0.26.1" [[test]] name = "test_tool_macros" required-features = ["server", "client"] @@ -298,11 +295,6 @@ 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"] @@ -348,15 +340,6 @@ 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" diff --git a/crates/rmcp/README.md b/crates/rmcp/README.md index f34469e..bb7837e 100644 --- a/crates/rmcp/README.md +++ b/crates/rmcp/README.md @@ -24,7 +24,6 @@ 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 @@ -46,10 +45,6 @@ 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/handler/client.rs b/crates/rmcp/src/handler/client.rs index 486fa3b..4143870 100644 --- a/crates/rmcp/src/handler/client.rs +++ b/crates/rmcp/src/handler/client.rs @@ -274,8 +274,8 @@ macro_rules! client_handler_methods { std::future::ready(()) } - fn get_info(&self) -> ClientConfig { - ClientConfig::default() + fn get_info(&self) -> ClientInfo { + ClientInfo::default() } }; } @@ -292,12 +292,12 @@ pub trait ClientHandler: Sized + 'static { client_handler_methods!(); } -/// Do nothing, with the default client config. +/// Do nothing, with default client info. impl ClientHandler for () {} -/// Do nothing, with a specific client config. -impl ClientHandler for ClientConfig { - fn get_info(&self) -> ClientConfig { +/// Do nothing, with a specific client info. +impl ClientHandler for ClientInfo { + fn get_info(&self) -> ClientInfo { self.clone() } } @@ -422,7 +422,7 @@ macro_rules! impl_client_handler_for_wrapper { (**self).on_custom_notification(notification, context) } - fn get_info(&self) -> ClientConfig { + fn get_info(&self) -> ClientInfo { (**self).get_info() } } diff --git a/crates/rmcp/src/handler/server.rs b/crates/rmcp/src/handler/server.rs index f846724..6dc7883 100644 --- a/crates/rmcp/src/handler/server.rs +++ b/crates/rmcp/src/handler/server.rs @@ -321,59 +321,13 @@ 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, - std::mem::take(&mut info.protocol_version), + info.protocol_version, &self.supported_protocol_versions(), - )?; - Ok(info) + ); + std::future::ready(Ok(info)) } /// Return the protocol versions supported by this server. /// @@ -382,10 +336,6 @@ macro_rules! server_handler_methods { /// 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) } @@ -600,8 +550,8 @@ macro_rules! server_handler_methods { std::future::ready(()) } - fn get_info(&self) -> ServerConfig { - ServerConfig::default() + fn get_info(&self) -> ServerInfo { + ServerInfo::default() } /// SEP-2663 `tasks/get`: return the current [`DetailedTask`] state. @@ -668,13 +618,6 @@ 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() } @@ -839,7 +782,7 @@ macro_rules! impl_server_handler_for_wrapper { (**self).on_custom_notification(notification, context) } - fn get_info(&self) -> ServerConfig { + fn get_info(&self) -> ServerInfo { (**self).get_info() } diff --git a/crates/rmcp/src/model.rs b/crates/rmcp/src/model.rs index ec7b93b..6a14098 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, oldest first. + /// All protocol versions known to this SDK. pub const KNOWN_VERSIONS: &[Self] = &[ Self::V_2024_11_05, Self::V_2025_03_26, @@ -190,43 +190,6 @@ 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 { @@ -1119,58 +1082,7 @@ 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. @@ -1330,18 +1242,18 @@ impl DiscoverResult { self } - /// Create a discovery result from the server's configuration. + /// Create a discovery result from the server's initialization information. pub fn from_server_info( supported_versions: Vec, - server_config: ServerConfig, + server_info: ServerInfo, ) -> Self { - let ServerConfig { + let ServerInfo { capabilities, server_info, instructions, meta, .. - } = server_config; + } = server_info; let mut result = Self { result_type: ResultType::COMPLETE, supported_versions, @@ -1383,9 +1295,9 @@ impl ServerPeerInfo { } #[allow(clippy::derivable_impls)] -impl Default for InitializeResult { +impl Default for ServerInfo { fn default() -> Self { - InitializeResult { + ServerInfo { protocol_version: ProtocolVersion::default(), capabilities: ServerCapabilities::default(), server_info: Implementation::from_build_env(), @@ -1396,9 +1308,9 @@ impl Default for InitializeResult { } #[allow(clippy::derivable_impls)] -impl Default for InitializeRequestParams { +impl Default for ClientInfo { fn default() -> Self { - InitializeRequestParams { + ClientInfo { meta: None, protocol_version: ProtocolVersion::default(), capabilities: ClientCapabilities::default(), @@ -4731,51 +4643,6 @@ mod tests { use super::*; - #[test] - fn known_versions_are_ordered_oldest_first() { - // `known_up_to` walks the list as a sorted prefix. - assert!( - ProtocolVersion::KNOWN_VERSIONS - .windows(2) - .all(|pair| pair[0].as_str() < pair[1].as_str()) - ); - } - - #[test] - fn known_up_to_includes_the_ceiling_itself() { - assert_eq!( - ProtocolVersion::known_up_to(&ProtocolVersion::V_2024_11_05), - &[ProtocolVersion::V_2024_11_05] - ); - } - - #[test] - fn known_up_to_the_newest_version_yields_every_known_version() { - assert_eq!( - ProtocolVersion::known_up_to(&ProtocolVersion::V_2026_07_28), - ProtocolVersion::KNOWN_VERSIONS - ); - } - - #[test] - fn known_up_to_an_unknown_ceiling_stops_at_the_versions_below_it() { - let unknown = ProtocolVersion(Cow::Borrowed("2025-07-01")); - assert_eq!( - ProtocolVersion::known_up_to(&unknown), - &[ - ProtocolVersion::V_2024_11_05, - ProtocolVersion::V_2025_03_26, - ProtocolVersion::V_2025_06_18, - ] - ); - } - - #[test] - fn known_up_to_a_ceiling_below_every_known_version_is_empty() { - let ancient = ProtocolVersion(Cow::Borrowed("1999-01-01")); - assert!(ProtocolVersion::known_up_to(&ancient).is_empty()); - } - #[cfg(feature = "transport-streamable-http-client")] #[test] fn transport_closed_marker_accepts_only_the_process_local_token() { diff --git a/crates/rmcp/src/model/elicitation_schema.rs b/crates/rmcp/src/model/elicitation_schema.rs index 4861087..78f1093 100644 --- a/crates/rmcp/src/model/elicitation_schema.rs +++ b/crates/rmcp/src/model/elicitation_schema.rs @@ -1116,15 +1116,6 @@ impl EnumSchema { #[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, @@ -1153,8 +1144,6 @@ pub struct ElicitationSchema { #[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")] @@ -1169,7 +1158,6 @@ struct ElicitationSchemaWire { 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()), @@ -1195,7 +1183,6 @@ impl From for ElicitationSchemaWire { properties.extend(remaining); Self { - schema: schema.schema, type_: schema.type_, title: schema.title, properties, @@ -1219,7 +1206,6 @@ impl ElicitationSchema { pub fn new(properties: BTreeMap) -> Self { let property_order = Some(properties.keys().cloned().collect()); Self { - schema: None, type_: ObjectTypeConst, title: None, properties, @@ -1315,12 +1301,6 @@ 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() @@ -1723,7 +1703,6 @@ impl ElicitationSchemaBuilder { let property_order = Some(self.properties.keys().cloned().collect()); Ok(ElicitationSchema { - schema: None, type_: ObjectTypeConst, title: self.title, properties: self.properties, @@ -1923,54 +1902,6 @@ mod tests { 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(()) - } - #[test] fn test_legacy_enum_schema_no_enum_names_omits_field() -> anyhow::Result<()> { // `LegacyEnumSchema` with `enum_names: None` must not serialize `"enumNames": null`. diff --git a/crates/rmcp/src/model/request_state.rs b/crates/rmcp/src/model/request_state.rs index 9f45572..5ed0f7d 100644 --- a/crates/rmcp/src/model/request_state.rs +++ b/crates/rmcp/src/model/request_state.rs @@ -40,10 +40,9 @@ //! //! ``` //! use rmcp::model::{RequestStateCodec, SealOptions}; -//! # fn main() -> Result<(), rmcp::model::RequestStateError> { //! //! // Derive the key from a per-process secret; keep it out of client reach. -//! let codec = RequestStateCodec::try_new(b"a-32-byte-or-longer-secret-key!!!")?; +//! let codec = RequestStateCodec::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"; @@ -54,60 +53,38 @@ //! //! // 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)?; +//! let opened = codec.open_with(&sealed, context).expect("integrity check passes"); //! assert_eq!(opened, b"step=2"); //! //! // A different principal (different context) is rejected. //! assert!(codec.open_with(&sealed, b"user:bob|tools/call:weather").is_err()); -//! # Ok(()) -//! # } //! ``` -use std::{ - collections::{HashMap, hash_map::Entry}, - time::Duration, -}; +use std::time::Duration; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; use hmac::{Hmac, KeyInit, Mac}; use serde::{Serialize, de::DeserializeOwned}; use sha2::Sha256; use thiserror::Error; -use zeroize::Zeroizing; type HmacSha256 = Hmac; -const VERSION_V1: &str = "rs1"; - -const VERSION_V2: &str = "rs2"; +/// Version tag prefixing every sealed value, so the wire format can evolve. +const VERSION: &str = "rs1"; /// Domain-separation label mixed into the HMAC so a `requestState` tag can never /// be confused with an HMAC computed for some other purpose using the same key. -const DOMAIN_V1: &[u8] = b"rmcp/mrtr/request-state/v1"; - -/// Domain-separation label for keyed request-state tags. -const DOMAIN_V2: &[u8] = b"rmcp/mrtr/request-state/v2"; +const DOMAIN: &[u8] = b"rmcp/mrtr/request-state/v1"; /// Length of the big-endian expiry prefix (unix milliseconds) stored at the /// front of every sealed body. `0` means "no expiry". const EXPIRY_LEN: usize = 8; -const MAX_KID_LEN: usize = 255; - -/// Maximum unpadded base64url length for [`MAX_KID_LEN`] bytes. -const MAX_ENCODED_KID_LEN: usize = 340; - -/// Errors returned when constructing or configuring a -/// [`RequestStateCodec`], or when processing a sealed value. +/// Errors returned when opening a sealed [`RequestStateCodec`] 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")] @@ -126,18 +103,6 @@ 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), @@ -183,261 +148,35 @@ impl<'a> SealOptions<'a> { } } -/// A keyed codec for integrity-protected SEP-2322 `requestState` values. +/// A keyed codec that seals and opens SEP-2322 `requestState` values with +/// HMAC-SHA256 integrity protection. /// -/// [`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. +/// 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). /// -/// 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(()) -/// # } -/// ``` +/// The key may be any length; HMAC internally normalizes it. For meaningful +/// security use a high-entropy key of at least 32 bytes. #[derive(Clone)] pub struct RequestStateCodec { - keys: Keys, -} - -#[derive(Clone)] -enum Keys { - Single(Zeroizing>), - Ring { - keys: HashMap>>, - seal_mode: SealMode, - rs1_fallbacks: Vec, - }, -} - -#[derive(Clone, Debug)] -enum SealMode { - Rs1 { key_id: String }, - Rs2 { key_id: String }, + key: Box<[u8]>, } impl std::fmt::Debug for RequestStateCodec { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - // Never leak signing or verification keys through Debug output. - match &self.keys { - Keys::Single(_) => f - .debug_struct("RequestStateCodec") - .field("mode", &"single") - .field("key", &"") - .finish(), - Keys::Ring { - keys, - seal_mode, - rs1_fallbacks, - } => f - .debug_struct("RequestStateCodec") - .field("mode", &"ring") - .field("key_count", &keys.len()) - .field("keys", &"") - .field("seal_mode", seal_mode) - .field("rs1_fallbacks", rs1_fallbacks) - .finish(), - } + // Never leak the signing key through Debug output. + f.debug_struct("RequestStateCodec") + .field("key", &"") + .finish() } } impl RequestStateCodec { - /// Minimum accepted signing-key length in bytes. - pub const MIN_KEY_LENGTH: usize = 32; - - /// Creates a legacy single-key codec that seals and opens `rs1` values - /// without validating the signing-key length. - #[deprecated( - since = "3.1.4", - note = "use RequestStateCodec::try_new to enforce the minimum signing-key length" - )] + /// Creates a codec from a signing key. pub fn new(key: impl Into>) -> Self { - Self::new_unchecked(key) - } - - /// Creates a codec without validating the signing-key length. - /// - /// Prefer [`try_new`](Self::try_new) unless the key length has already been - /// validated. - pub fn new_unchecked(key: impl Into>) -> Self { Self { - keys: Keys::Single(Zeroizing::new(key.into())), - } - } - - /// Creates a legacy single-key codec after validating its signing-key length. - /// - /// # Errors - /// - /// Returns [`RequestStateError::KeyTooShort`] when `key` contains fewer - /// than [`MIN_KEY_LENGTH`](Self::MIN_KEY_LENGTH) bytes. - pub fn try_new(key: impl Into>) -> Result { - let key = Zeroizing::new(key.into()); - Self::validate_key_length(&key)?; - Ok(Self { - keys: Keys::Single(key), - }) - } - - /// Creates a keyring codec that seals `rs2` with `active_kid`. - /// - /// All configured keys can open matching `rs2` values. Legacy `rs1` values - /// require [`with_rs1_fallback`](Self::with_rs1_fallback) or - /// [`with_rs1_signing`](Self::with_rs1_signing). - /// Key ids are case-sensitive UTF-8 strings of 1 to 255 bytes. - /// - /// # Errors - /// - /// Returns [`RequestStateError::KeyTooShort`] if a configured key contains - /// fewer than [`MIN_KEY_LENGTH`](Self::MIN_KEY_LENGTH) bytes. Returns - /// [`RequestStateError::InvalidKeyring`] if the keyring is empty, contains - /// duplicate or invalid key ids, or does not contain `active_kid`. - pub fn new_with_keyring( - active_kid: impl Into, - keys: impl IntoIterator, - ) -> Result - where - K: Into, - V: Into>, - { - let mut ring = HashMap::new(); - for (kid, key) in keys { - let kid = kid.into(); - let key = Zeroizing::new(key.into()); - Self::validate_config_kid(&kid)?; - Self::validate_key_length(&key)?; - match ring.entry(kid) { - Entry::Vacant(entry) => { - entry.insert(key); - } - Entry::Occupied(_) => { - return Err(RequestStateError::InvalidKeyring("duplicate key id")); - } - } - } - - if ring.is_empty() { - return Err(RequestStateError::InvalidKeyring( - "keyring must contain at least one key", - )); - } - - let active_kid = active_kid.into(); - Self::validate_config_kid(&active_kid)?; - if !ring.contains_key(&active_kid) { - return Err(RequestStateError::InvalidKeyring( - "active key id is not present in the keyring", - )); - } - - Ok(Self { - keys: Keys::Ring { - keys: ring, - seal_mode: SealMode::Rs2 { key_id: active_kid }, - rs1_fallbacks: Vec::new(), - }, - }) - } - - /// Uses `kid` to sign legacy `rs1` values during a rolling migration. - /// The codec accepts its own `rs1` output and configured `rs2` values. - /// - /// # Errors - /// - /// Returns [`RequestStateError::InvalidKeyring`] when called on a single-key - /// codec, or if the keyring does not contain `kid`. - pub fn with_rs1_signing(mut self, kid: impl AsRef) -> Result { - let kid = kid.as_ref(); - match &mut self.keys { - Keys::Single(_) => Err(RequestStateError::InvalidKeyring( - "rs1 transitional signing requires a keyring", - )), - Keys::Ring { - keys, - seal_mode, - rs1_fallbacks, - } => { - if !keys.contains_key(kid) { - return Err(RequestStateError::InvalidKeyring( - "rs1 signing key id is not present in the keyring", - )); - } - - let kid = kid.to_owned(); - *seal_mode = SealMode::Rs1 { - key_id: kid.clone(), - }; - if !rs1_fallbacks.iter().any(|existing| existing == &kid) { - rs1_fallbacks.push(kid); - } - Ok(self) - } - } - } - - /// Accepts legacy `rs1` values signed with `kid`. - /// Adding the same id more than once has no effect; remove fallbacks after - /// old values have drained. - /// - /// # Errors - /// - /// Returns [`RequestStateError::InvalidKeyring`] when called on a single-key - /// codec, or if the keyring does not contain `kid`. - pub fn with_rs1_fallback(mut self, kid: impl AsRef) -> Result { - let kid = kid.as_ref(); - match &mut self.keys { - Keys::Single(_) => Err(RequestStateError::InvalidKeyring( - "rs1 fallbacks require a keyring", - )), - Keys::Ring { - keys, - rs1_fallbacks, - .. - } => { - if !keys.contains_key(kid) { - return Err(RequestStateError::InvalidKeyring( - "rs1 fallback key id is not present in the keyring", - )); - } - if !rs1_fallbacks.iter().any(|existing| existing == kid) { - rs1_fallbacks.push(kid.to_owned()); - } - Ok(self) - } + key: key.into().into_boxed_slice(), } } @@ -496,12 +235,12 @@ impl RequestStateCodec { /// /// # Errors /// - /// Returns a [`RequestStateError`] if the value is malformed, cannot be - /// verified, names an unavailable key, or has expired. - /// - /// Applications MUST map all token-opening failures to a single - /// client-facing error and reserve the detailed variants for internal - /// diagnostics. + /// - [`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. pub fn open_with( &self, sealed: &str, @@ -548,24 +287,21 @@ impl RequestStateCodec { body.extend_from_slice(&expiry.to_be_bytes()); body.extend_from_slice(payload); - match &self.keys { - Keys::Single(key) => Self::seal_rs1(key, options.associated_data, &body), - Keys::Ring { - keys, seal_mode, .. - } => match seal_mode { - SealMode::Rs1 { key_id } => Self::seal_rs1( - keys.get(key_id).expect("validated rs1 signing key"), - options.associated_data, - &body, - ), - SealMode::Rs2 { key_id } => Self::seal_rs2( - key_id, - keys.get(key_id).expect("validated rs2 signing key"), - options.associated_data, - &body, - ), - }, - } + let tag = self + .mac_for(options.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); + 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_at( @@ -573,71 +309,15 @@ impl RequestStateCodec { sealed: &str, associated_data: &[u8], now_ms: i64, - ) -> Result, RequestStateError> { - match sealed.split('.').next() { - Some(VERSION_V1) => self.open_rs1_at(sealed, associated_data, now_ms), - Some(VERSION_V2) => self.open_rs2_at(sealed, associated_data, now_ms), - _ => Err(RequestStateError::MalformedFormat), - } - } - - fn seal_rs1(key: &[u8], associated_data: &[u8], body: &[u8]) -> String { - let tag = Self::mac_v1(key, associated_data, body) - .finalize() - .into_bytes(); - let mut out = String::with_capacity( - VERSION_V1.len() + 2 + Self::b64_len(body.len()) + Self::b64_len(tag.len()), - ); - out.push_str(VERSION_V1); - out.push('.'); - URL_SAFE_NO_PAD.encode_string(body, &mut out); - out.push('.'); - URL_SAFE_NO_PAD.encode_string(tag.as_slice(), &mut out); - out - } - - fn seal_rs2(kid: &str, key: &[u8], associated_data: &[u8], body: &[u8]) -> String { - let tag = Self::mac_v2(key, kid.as_bytes(), associated_data, body) - .finalize() - .into_bytes(); - let mut out = String::with_capacity( - VERSION_V2.len() - + 3 - + Self::b64_len(kid.len()) - + Self::b64_len(body.len()) - + Self::b64_len(tag.len()), - ); - out.push_str(VERSION_V2); - out.push('.'); - URL_SAFE_NO_PAD.encode_string(kid.as_bytes(), &mut out); - out.push('.'); - URL_SAFE_NO_PAD.encode_string(body, &mut out); - out.push('.'); - URL_SAFE_NO_PAD.encode_string(tag.as_slice(), &mut out); - out - } - - fn open_rs1_at( - &self, - sealed: &str, - associated_data: &[u8], - now_ms: i64, ) -> Result, RequestStateError> { let mut parts = sealed.split('.'); let version = parts.next().ok_or(RequestStateError::MalformedFormat)?; let body_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?; let tag_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?; - if parts.next().is_some() || version != VERSION_V1 { + if parts.next().is_some() || version != VERSION { 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)?; @@ -645,81 +325,11 @@ impl RequestStateCodec { .decode(tag_b64) .map_err(|_| RequestStateError::InvalidEncoding)?; - // Authentication must not short-circuit based on tag bytes. - match &self.keys { - Keys::Single(key) => Self::mac_v1(key, associated_data, &body) - .verify_slice(&tag) - .map_err(|_| RequestStateError::IntegrityCheckFailed)?, - Keys::Ring { - keys, - rs1_fallbacks, - .. - } => { - // Try every fallback so timing does not identify the matching key. - let mut verified = false; - for kid in rs1_fallbacks { - let key = keys.get(kid).expect("validated rs1 fallback key"); - let matches = Self::mac_v1(key, associated_data, &body) - .verify_slice(&tag) - .is_ok(); - verified |= matches; - } - if !verified { - return Err(RequestStateError::IntegrityCheckFailed); - } - } - } - - Self::open_authenticated_body(body, now_ms) - } - - fn open_rs2_at( - &self, - sealed: &str, - associated_data: &[u8], - now_ms: i64, - ) -> Result, RequestStateError> { - let mut parts = sealed.split('.'); - let version = parts.next().ok_or(RequestStateError::MalformedFormat)?; - let kid_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?; - let body_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?; - let tag_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?; - if parts.next().is_some() || version != VERSION_V2 { - return Err(RequestStateError::MalformedFormat); - } - - if kid_b64.len() > MAX_ENCODED_KID_LEN { - return Err(RequestStateError::InvalidKeyId); - } - let kid = URL_SAFE_NO_PAD - .decode(kid_b64) - .map_err(|_| RequestStateError::InvalidEncoding)?; - if kid.is_empty() || kid.len() > MAX_KID_LEN { - return Err(RequestStateError::InvalidKeyId); - } - let kid = String::from_utf8(kid).map_err(|_| RequestStateError::InvalidKeyId)?; - - let key = match &self.keys { - Keys::Single(_) => return Err(RequestStateError::UnknownKeyId), - Keys::Ring { keys, .. } => keys.get(&kid).ok_or(RequestStateError::UnknownKeyId)?, - }; - - let body = URL_SAFE_NO_PAD - .decode(body_b64) - .map_err(|_| RequestStateError::InvalidEncoding)?; - let tag = URL_SAFE_NO_PAD - .decode(tag_b64) - .map_err(|_| RequestStateError::InvalidEncoding)?; - - // Authentication must not short-circuit based on tag bytes. - Self::mac_v2(key, kid.as_bytes(), associated_data, &body) + // `verify_slice` compares in constant time and rejects wrong-length tags. + self.mac_for(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); @@ -732,55 +342,20 @@ impl RequestStateCodec { Ok(body[EXPIRY_LEN..].to_vec()) } - fn mac_v1(key: &[u8], associated_data: &[u8], body: &[u8]) -> HmacSha256 { - let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts keys of any length"); - mac.update(DOMAIN_V1); + /// 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); 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() } @@ -790,42 +365,9 @@ 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::try_new(b"test-key-test-key-test-key-32byte".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"test-key-test-key-test-key-32byte".to_vec()); let sealed = codec.seal(b"hello world"); assert!(sealed.starts_with("rs1.")); assert_eq!(codec.open(&sealed).unwrap(), b"hello world"); @@ -838,8 +380,7 @@ mod tests { tool: String, round: u32, } - let codec = - RequestStateCodec::try_new(b"another-strong-signing-key-here!!".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"another-strong-signing-key-here!!".to_vec()); let state = State { tool: "weather".into(), round: 3, @@ -851,16 +392,14 @@ mod tests { #[test] fn empty_payload_roundtrips() { - let codec = - RequestStateCodec::try_new(b"empty-payload-test-signing-key!!".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"k".to_vec()); let sealed = codec.seal(b""); assert_eq!(codec.open(&sealed).unwrap(), b""); } #[test] fn tampered_payload_is_rejected() { - let codec = - RequestStateCodec::try_new(b"signing-key-signing-key-signing!!".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"signing-key-signing-key-signing!!".to_vec()); let sealed = codec.seal(b"amount=100"); // Replace the body section but keep the original tag. @@ -877,10 +416,8 @@ mod tests { #[test] fn different_key_is_rejected() { - let signer = - RequestStateCodec::try_new(b"the-real-signing-key-value-here!!".to_vec()).unwrap(); - let attacker = - RequestStateCodec::try_new(b"a-totally-different-forged-key!!!".to_vec()).unwrap(); + let 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 sealed = signer.seal(b"trusted"); assert!(matches!( attacker.open(&sealed), @@ -890,19 +427,17 @@ mod tests { #[test] fn appended_bytes_are_rejected() { - let codec = - RequestStateCodec::try_new(b"key-key-key-key-key-key-key-key!!".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec()); let mut sealed = codec.seal(b"state"); sealed.push('x'); assert!(codec.open(&sealed).is_err()); } #[test] - fn unsupported_version_is_malformed() { - let codec = - RequestStateCodec::try_new(b"wrong-version-test-signing-key!!".to_vec()).unwrap(); + fn wrong_version_prefix_is_malformed() { + let codec = RequestStateCodec::new(b"key".to_vec()); let sealed = codec.seal(b"state"); - let bumped = sealed.replacen("rs1.", "rs3.", 1); + let bumped = sealed.replacen("rs1.", "rs2.", 1); assert!(matches!( codec.open(&bumped), Err(RequestStateError::MalformedFormat) @@ -911,8 +446,7 @@ mod tests { #[test] fn missing_sections_are_malformed() { - let codec = - RequestStateCodec::try_new(b"missing-sections-test-signing-key".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"key".to_vec()); assert!(matches!( codec.open("rs1"), Err(RequestStateError::MalformedFormat) @@ -929,8 +463,7 @@ mod tests { #[test] fn non_base64_sections_are_invalid_encoding() { - let codec = - RequestStateCodec::try_new(b"invalid-base64-test-signing-key!!".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"key".to_vec()); assert!(matches!( codec.open("rs1.!!!!.!!!!"), Err(RequestStateError::InvalidEncoding) @@ -939,8 +472,7 @@ mod tests { #[test] fn debug_does_not_leak_key() { - let codec = - RequestStateCodec::try_new(b"super-secret-key-super-secret-key!!".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"super-secret-key".to_vec()); let rendered = format!("{codec:?}"); assert!(!rendered.contains("super-secret-key")); assert!(rendered.contains("redacted")); @@ -951,8 +483,7 @@ mod tests { #[test] fn matching_context_opens() { - let codec = - RequestStateCodec::try_new(b"key-key-key-key-key-key-key-key!!".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec()); 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"); @@ -960,8 +491,7 @@ mod tests { #[test] fn different_context_is_rejected() { - let codec = - RequestStateCodec::try_new(b"key-key-key-key-key-key-key-key!!".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec()); let sealed = codec.seal_with(b"state", &SealOptions::new().associated_data(b"user:alice")); assert!(matches!( @@ -972,8 +502,7 @@ mod tests { #[test] fn missing_context_is_rejected() { - let codec = - RequestStateCodec::try_new(b"key-key-key-key-key-key-key-key!!".to_vec()).unwrap(); + let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec()); let sealed = codec.seal_with(b"state", &SealOptions::new().associated_data(b"user:alice")); // Opening without the associated data must fail closed. @@ -991,7 +520,7 @@ mod tests { #[test] fn within_ttl_opens() { - let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); + let codec = RequestStateCodec::new(KEY.to_vec()); let sealed = codec.seal_at( b"state", &SealOptions::new().ttl(Duration::from_secs(60)), @@ -1003,7 +532,7 @@ mod tests { #[test] fn past_ttl_is_expired() { - let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); + let codec = RequestStateCodec::new(KEY.to_vec()); let sealed = codec.seal_at( b"state", &SealOptions::new().ttl(Duration::from_secs(60)), @@ -1018,14 +547,14 @@ mod tests { #[test] fn no_ttl_never_expires() { - let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); + let codec = RequestStateCodec::new(KEY.to_vec()); let sealed = codec.seal_at(b"state", &SealOptions::new(), 1_000); assert_eq!(codec.open_at(&sealed, &[], i64::MAX).unwrap(), b"state"); } #[test] fn ttl_and_associated_data_combine() { - let codec = RequestStateCodec::try_new(KEY.to_vec()).unwrap(); + let codec = RequestStateCodec::new(KEY.to_vec()); let ctx = b"user:alice"; let sealed = codec.seal_at( b"state", @@ -1045,450 +574,4 @@ 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/service.rs b/crates/rmcp/src/service.rs index 63b38e4..20fd2e9 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -136,14 +136,6 @@ 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 @@ -201,17 +193,12 @@ pub enum PeerRequestAssociation { Unknown { has_pending_outbound_request: bool }, } -/// Whether `version` predates `2026-07-28`, the revision that replaced the -/// `initialize` handshake with per-request metadata. -pub(crate) fn is_legacy_version(version: &ProtocolVersion) -> bool { - version.as_str() < ProtocolVersion::V_2026_07_28.as_str() -} - pub(crate) fn uses_legacy_lifecycle( protocol_version: Option<&ProtocolVersion>, uses_discover_lifecycle: bool, ) -> bool { - !uses_discover_lifecycle && protocol_version.is_none_or(is_legacy_version) + !uses_discover_lifecycle + && protocol_version.is_none_or(|version| version < &ProtocolVersion::V_2026_07_28) } pub(crate) fn peer_request_association( @@ -530,10 +517,6 @@ 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>; @@ -1310,14 +1293,7 @@ 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, - VecDeque::new(), - ct, - ) + serve_inner(service, transport.into_transport(), peer, peer_rx, ct) } /// Spawn a task that may hold `!Send` state when the `local` feature is active. @@ -1342,19 +1318,12 @@ 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 @@ -1374,7 +1343,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 @@ -1387,7 +1356,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 = initial_messages; + let mut batch_messages = VecDeque::>::new(); let mut send_task_set = tokio::task::JoinSet::::new(); let mut response_send_tasks = tokio::task::JoinSet::<()>::new(); #[derive(Debug)] @@ -1471,7 +1440,7 @@ where Event::SendTaskResult(SendTaskResult::Request { id, result }) => { if let Err(e) = result && let Some(responder) = local_responder_pool.remove(&id) { - let _ = responder.responder.send(Err(ServiceError::TransportSend(e))); + let _ = responder.send(Err(ServiceError::TransportSend(e))); } } Event::SendTaskResult(SendTaskResult::Notification { @@ -1489,7 +1458,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.responder.send(Err(ServiceError::Cancelled { + let _response_result = responder.send(Err(ServiceError::Cancelled { reason: param.reason.clone(), })); } @@ -1526,10 +1495,7 @@ where id, responder, }) => { - local_responder_pool.insert(id.clone(), PendingResponse { - responder, - is_subscription: R::is_subscription_request(&request), - }); + local_responder_pool.insert(id.clone(), responder); let send = transport.send(JsonRpcMessage::request(request, id.clone())); { let id = id.clone(); @@ -1630,38 +1596,31 @@ where })) => { tracing::info!(?notification, "received notification"); R::invalidate_response_cache(&peer, ¬ification).await; - let subscription_id = + let cancellation_request_id = if let Some(cancelled) = R::peer_cancelled_params(¬ification) { let request_id = cancelled.request_id.clone(); - // Modern servers cancel listen requests; legacy peers cancel - // requests they originated, even when both directions share an ID. - if R::peer_cancels_subscriptions(&peer) { - if let Some(request_id) = request_id.as_ref() { - if local_responder_pool.get(request_id) - .is_some_and(|pending| pending.is_subscription) - && let Some(pending) = local_responder_pool.remove(request_id) + if let Some(request_id) = request_id.as_ref() { + if R::IS_CLIENT { + if let Some(responder) = + local_responder_pool.remove(request_id) { - let _ = pending.responder.send(Err(ServiceError::Cancelled { + let _ = responder.send(Err(ServiceError::Cancelled { reason: cancelled.reason.clone(), })); - } else { - tracing::debug!(%request_id, "ignoring cancellation of unknown subscription"); - continue; } - } - request_id - } else { - if let Some(request_id) = request_id.as_ref() - && let Some(ct) = local_ct_pool.remove(request_id) - { + } else if let Some(ct) = local_ct_pool.remove(request_id) { tracing::info!(id = %request_id, reason = cancelled.reason, "cancelled"); ct.cancel(); } - None } + request_id } else { - notification.get_meta().subscription_id() + None }; + 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) @@ -1678,7 +1637,7 @@ where && let Some(responder) = local_responder_pool.remove(&subscription_id) { - let _ = responder.responder + let _ = responder .send(Err(ServiceError::SubscriptionLagged { capacity })); } peer.unregister_subscription(&subscription_id); @@ -1730,7 +1689,7 @@ where if let Some(responder) = remove_pending_request(&mut local_responder_pool, &id) { - let response_result = responder.responder.send(Ok(result)); + let response_result = responder.send(Ok(result)); if let Err(_error) = response_result { tracing::warn!(%id, "Error sending response"); } @@ -1751,7 +1710,7 @@ where } else { ServiceError::McpError(error) }; - let _response_result = responder.responder.send(Err(service_error)); + let _response_result = responder.send(Err(service_error)); if let Err(_error) = _response_result { tracing::warn!(%id, "Error sending response"); } @@ -1785,17 +1744,6 @@ 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; diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 7d0d794..520410f 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -13,13 +13,13 @@ use crate::{ model::{ ArgumentInfo, CacheScope, CallToolRequest, CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, CancelTaskRequest, CancelledNotification, - CancelledNotificationParam, ClientConfig, ClientJsonRpcMessage, ClientNotification, + CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage, ClientNotification, ClientRequest, ClientResult, CompleteRequest, CompleteRequestParams, CompleteResult, CompletionContext, CompletionInfo, DEFAULT_MRTR_MAX_ROUNDS, DiscoverRequest, DiscoverRequestParams, DiscoverResult, ErrorData, GetExtensions, GetMeta, GetPromptRequest, GetPromptRequestParams, GetPromptResponse, GetPromptResult, GetTaskParams, GetTaskRequest, - GetTaskResult, InitializeRequest, InitializeRequestParams, InitializedNotification, - InputRequest, InputRequiredResult, InputResponses, JsonRpcResponse, ListPromptsRequest, + GetTaskResult, InitializeRequest, InitializedNotification, InputRequest, + InputRequiredResult, InputResponses, JsonRpcResponse, ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest, ListResourceTemplatesResult, ListResourcesRequest, ListResourcesResult, ListToolsRequest, ListToolsResult, NumberOrString, PaginatedRequestParams, ProgressNotification, ProgressNotificationParam, @@ -266,7 +266,7 @@ impl ServiceRole for RoleClient { type PeerReq = ServerRequest; type PeerResp = ServerResult; type PeerNot = ServerNotification; - type Info = ClientConfig; + type Info = ClientInfo; type PeerInfo = ServerPeerInfo; type InitializeError = ClientInitializeError; const IS_CLIENT: bool = true; @@ -292,20 +292,6 @@ impl ServiceRole for RoleClient { } } - 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. @@ -841,14 +827,7 @@ where } } } - Ok(serve_inner( - service, - transport, - peer, - peer_rx, - VecDeque::new(), - ct, - )) + Ok(serve_inner(service, transport, peer, peer_rx, ct)) } /// Modern-era JSON-RPC error codes a server can return from `server/discover` @@ -891,7 +870,7 @@ async fn legacy_startup( transport: &mut T, id_provider: &Arc, peer: &Peer, - client_info: InitializeRequestParams, + client_info: ClientInfo, ) -> Result<(), ClientInitializeError> where S: Service, @@ -940,7 +919,7 @@ async fn discover_startup( transport: &mut T, id_provider: &Arc, peer: &Peer, - client_info: &InitializeRequestParams, + client_info: &ClientInfo, preferred_versions: Vec, ) -> Result where @@ -2214,232 +2193,6 @@ 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}; diff --git a/crates/rmcp/src/service/server.rs b/crates/rmcp/src/service/server.rs index 3ff229f..8efbd34 100644 --- a/crates/rmcp/src/service/server.rs +++ b/crates/rmcp/src/service/server.rs @@ -13,16 +13,15 @@ use super::*; use crate::model::{ElicitRequest, ElicitRequestParams, ElicitResult, ElicitationAction}; use crate::{ model::{ - CancelledNotification, CancelledNotificationParam, ClientJsonRpcMessage, + CancelledNotification, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage, ClientNotification, ClientRequest, ClientResult, CreateMessageRequest, - CreateMessageRequestParams, CreateMessageResult, EmptyResult, ErrorData, - InitializeRequestParams, ListRootsRequest, ListRootsResult, LoggingMessageNotification, - LoggingMessageNotificationParam, ProgressNotification, ProgressNotificationParam, - PromptListChangedNotification, ProtocolVersion, ResourceListChangedNotification, - ResourceUpdatedNotification, ResourceUpdatedNotificationParam, ServerConfig, - ServerNotification, ServerRequest, ServerResult, SubscriptionFilter, - SubscriptionsAcknowledgedNotification, SubscriptionsAcknowledgedNotificationParams, - ToolListChangedNotification, + CreateMessageRequestParams, CreateMessageResult, EmptyResult, ErrorData, ListRootsRequest, + ListRootsResult, LoggingMessageNotification, LoggingMessageNotificationParam, + ProgressNotification, ProgressNotificationParam, PromptListChangedNotification, + ProtocolVersion, ResourceListChangedNotification, ResourceUpdatedNotification, + ResourceUpdatedNotificationParam, ServerInfo, ServerNotification, ServerRequest, + ServerResult, SubscriptionFilter, SubscriptionsAcknowledgedNotification, + SubscriptionsAcknowledgedNotificationParams, ToolListChangedNotification, }, transport::DynamicTransportError, }; @@ -38,8 +37,8 @@ impl ServiceRole for RoleServer { type PeerReq = ClientRequest; type PeerResp = ClientResult; type PeerNot = ClientNotification; - type Info = ServerConfig; - type PeerInfo = InitializeRequestParams; + type Info = ServerInfo; + type PeerInfo = ClientInfo; type InitializeError = ServerInitializeError; const IS_CLIENT: bool = false; @@ -461,92 +460,26 @@ where } } -/// Echoes the client-requested version if the server can serve it over the -/// `initialize` handshake; otherwise returns a legacy version the server does -/// support. +/// Echoes the client-requested version if the server supports it; otherwise +/// returns `server_fallback`. /// /// `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 +/// version it cannot serve. pub(crate) fn negotiate_protocol_version( client_requested: &ProtocolVersion, server_fallback: ProtocolVersion, 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) +) -> ProtocolVersion { + if server_supported.contains(client_requested) { + client_requested.clone() } else { - newest_legacy_version(server_supported) - }; - let Some(legacy_fallback) = legacy_fallback else { tracing::warn!( client_requested = %client_requested, - "server supports no protocol version with an initialize handshake; rejecting" + server_fallback = %server_fallback, + "client requested unsupported protocol version; falling back to server default" ); - 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") - } + server_fallback } } @@ -594,40 +527,33 @@ where let initialize_request = match request { ClientRequest::InitializeRequest(request) => request, - request => { - let missing_metadata = request + mut request => { + if !request .get_meta() - .missing_required_keys(&ProtocolVersion::V_2026_07_28); - if !missing_metadata.is_empty() { - transport - .send(ServerJsonRpcMessage::error( - missing_request_metadata_error(&missing_metadata), - Some(id.clone()), - )) - .await - .map_err(|error| { - ServerInitializeError::transport::( - error, - "sending pre-init metadata error response", - ) - })?; + .missing_required_keys(&ProtocolVersion::V_2026_07_28) + .is_empty() + { return Err(ServerInitializeError::ExpectedInitializeRequest(Some( ClientJsonRpcMessage::request(request, id), ))); } let (peer, peer_rx) = Peer::new(id_provider, None); peer.require_request_metadata(); - // Dispatch the request from inside the service loop rather than - // inline: its handler may send notifications through `peer`, which - // only complete once the loop drains `peer_rx`. - return Ok(serve_inner( - service, - transport, - peer, - peer_rx, - VecDeque::from([ClientJsonRpcMessage::request(request, id)]), - ct, - )); + let 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)); } }; let requested_protocol_version = initialize_request.params.protocol_version.clone(); @@ -642,21 +568,27 @@ where peer: peer.clone(), }; // Send initialize response - let mut init_response = match service.handle_request(request, context).await { + let init_response = service.handle_request(request, context).await; + let mut init_response = match init_response { Ok(ServerResult::InitializeResult(init_response)) => init_response, Ok(result) => { return Err(ServerInitializeError::UnexpectedInitializeResponse(result)); } - Err(e) => return Err(report_initialize_failure(&mut transport, e, id).await), + 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)); + } }; - init_response.protocol_version = match negotiate_protocol_version( + init_response.protocol_version = negotiate_protocol_version( &requested_protocol_version, init_response.protocol_version, &service.supported_protocol_versions(), - ) { - Ok(version) => version, - Err(e) => return Err(report_initialize_failure(&mut transport, e, id).await), - }; + ); // Update peer_info so context.protocol_version() reflects the negotiated // version in all subsequent request handlers. negotiated_peer_info.protocol_version = init_response.protocol_version.clone(); @@ -677,14 +609,7 @@ 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, - VecDeque::new(), - ct, - )) + Ok(serve_inner(service, transport, peer, peer_rx, ct)) } macro_rules! method { diff --git a/crates/rmcp/src/transport.rs b/crates/rmcp/src/transport.rs index fdedc9c..06fde8e 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, CredentialRefreshGuard, CredentialStore, + AuthorizedHttpClient, ClientCredentialsConfig, CredentialStore, EXTENSION_OAUTH_CLIENT_CREDENTIALS, InMemoryCredentialStore, InMemoryStateStore, OAuthHttpClient, OAuthHttpClientError, OAuthHttpClientFuture, OAuthHttpRedirectPolicy, OAuthHttpRequest, ScopeUpgradeConfig, StateStore, StoredAuthorizationState, StoredCredentials, diff --git a/crates/rmcp/src/transport/auth.rs b/crates/rmcp/src/transport/auth.rs index 0727368..02d5063 100644 --- a/crates/rmcp/src/transport/auth.rs +++ b/crates/rmcp/src/transport/auth.rs @@ -1,5 +1,5 @@ use std::{ - collections::{HashMap, HashSet}, + collections::HashMap, future::Future, net::{IpAddr, Ipv4Addr, Ipv6Addr}, pin::Pin, @@ -28,9 +28,6 @@ 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; @@ -102,19 +99,6 @@ 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, @@ -271,32 +255,11 @@ 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>; @@ -304,16 +267,6 @@ 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) @@ -562,9 +515,6 @@ 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), @@ -685,19 +635,6 @@ 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] @@ -977,7 +914,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(Clone)] +#[derive(Debug, Clone)] #[non_exhaustive] pub enum ClientCredentialsConfig { /// Client secret authentication (credentials in request body) @@ -1000,42 +937,6 @@ 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, @@ -1337,9 +1238,13 @@ 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(default_oauth_http_client()?), + Arc::new(ReqwestOAuthHttpClient::new(http_client)?), OAuthHttpRedirectPolicy::Stop, ) .await @@ -2262,21 +2167,8 @@ 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)?; @@ -2293,7 +2185,6 @@ impl AuthorizationManager { .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)); } @@ -2319,12 +2210,9 @@ impl AuthorizationManager { token_result.set_refresh_token(Some(refresh_token_value)); } - 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) + 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(), }; *self.current_scopes.write().await = granted_scopes.clone(); @@ -2538,50 +2426,26 @@ impl AuthorizationManager { async fn discover_oauth_server_via_resource_metadata( &self, ) -> Result, AuthError> { - let Some((resource_metadata_url, resource_metadata)) = - self.discover_resource_metadata().await? - else { + let Some(resource_metadata_url) = self.discover_resource_metadata_url().await? else { return Ok(None); }; - self.authorization_metadata_from_resource_metadata( - &resource_metadata_url, - resource_metadata, - ) - .await + self.discover_oauth_server_from_resource_metadata_url(&resource_metadata_url) + .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, - ResourceMetadataUrlOrigin::Advertised, - ) + .fetch_resource_metadata_from_url(resource_metadata_url) .await? else { return Ok(None); }; - self.authorization_metadata_from_resource_metadata(resource_metadata_url, resource_metadata) - .await - } + self.validate_resource_metadata_resource(&resource_metadata)?; - /// 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 @@ -2594,32 +2458,21 @@ impl AuthorizationManager { *self.resource_scopes.write().await = scopes; } - // 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()); - } - }; + let mut candidates = Vec::new(); if let Some(single) = resource_metadata.authorization_server { - push_candidate(single); + candidates.push(single); } - for candidate in resource_metadata.authorization_servers.unwrap_or_default() { - push_candidate(candidate); + if let Some(list) = resource_metadata.authorization_servers { + candidates.extend(list); } - 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 = candidate.trim(); + if candidate.is_empty() { + continue; + } - for candidate in &candidates { let candidate_url = match Url::parse(candidate) { Ok(url) => url, Err(_) => match resource_metadata_url.join(candidate) { @@ -2656,12 +2509,7 @@ impl AuthorizationManager { } } - // 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(", ") - ))) + Ok(None) } fn validate_resource_metadata_resource( @@ -2721,115 +2569,52 @@ impl AuthorizationManager { || 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. + async fn discover_resource_metadata_url(&self) -> Result, AuthError> { + if let Some(resource_metadata_url) = + self.probe_resource_metadata_url(&self.base_url).await? + { + return Ok(Some(resource_metadata_url)); } - // The other place the document can be is the well-known location. + // If the primary URL doesn't use WWW-Authenticate, try oauth-protected-resource discovery. // 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 candidate_url = self.base_url.clone(); - candidate_url.set_query(None); - candidate_url.set_fragment(None); - candidate_url.set_path(&candidate_path); - - if !requested.insert(candidate_url.clone()) { - continue; - } - - let response = self - .discovery_get(&candidate_url) - .await - .map_err(|error| Self::discovery_failed(&candidate_url, error))?; - - match response.status() { - // The candidate url is the document itself, so read the body here - // instead of requesting the same url again. - StatusCode::OK => { - if let Some(metadata) = self.read_resource_metadata( - &candidate_url, - response.body(), - ResourceMetadataUrlOrigin::WellKnownGuess, - )? { - return Ok(Some((candidate_url, metadata))); - } - } - StatusCode::UNAUTHORIZED => { - let Some(pointer_url) = self - .extract_resource_metadata_url_from_www_authenticate(&response) - .await - else { - continue; - }; - if !requested.insert(pointer_url.clone()) { - continue; - } - // Found on a url this run guessed at, so a wrong answer only rules - // out the pointer. - if let Some(metadata) = self - .fetch_resource_metadata_from_url( - &pointer_url, - ResourceMetadataUrlOrigin::WellKnownGuess, - ) - .await? - { - return Ok(Some((pointer_url, metadata))); - } - } - status => debug!("resource metadata probe returned unexpected status: {status}"), + 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 Some(resource_metadata_url) = + self.probe_resource_metadata_url(&discovery_url).await? + { + return Ok(Some(resource_metadata_url)); } } Ok(None) } - /// 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. + /// Probe `url` with a GET, extracting the resource metadata url from a + /// 200 (the url itself is the metadata document) or from a 401's + /// WWW-Authenticate header value. /// https://www.rfc-editor.org/rfc/rfc9728.html#name-use-of-www-authenticate-for - async fn probe_resource_endpoint_for_challenge(&self) -> Result, AuthError> { + async fn probe_resource_metadata_url(&self, url: &Url) -> Result, AuthError> { let response = self - .discovery_get(&self.base_url) + .discovery_get(url) .await - .map_err(|error| Self::discovery_failed(&self.base_url, error))?; + .map_err(|error| Self::discovery_failed(url, error))?; - if response.status() == StatusCode::UNAUTHORIZED { - return Ok(self + match response.status() { + StatusCode::OK => Ok(Some(url.clone())), + StatusCode::UNAUTHORIZED => Ok(self .extract_resource_metadata_url_from_www_authenticate(&response) - .await); + .await), + status => { + debug!("resource metadata probe returned unexpected status: {status}"); + Ok(None) + } } - - debug!( - "resource endpoint probe returned {}, no WWW-Authenticate pointer to follow", - response.status() - ); - Ok(None) } async fn extract_resource_metadata_url_from_www_authenticate( @@ -2855,7 +2640,6 @@ impl AuthorizationManager { async fn fetch_resource_metadata_from_url( &self, resource_metadata_url: &Url, - origin: ResourceMetadataUrlOrigin, ) -> Result, AuthError> { debug!( "resource metadata discovery url: {:?}", @@ -2874,69 +2658,13 @@ impl AuthorizationManager { return Ok(None); } - 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) { + let metadata = match serde_json::from_slice::(response.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)) } @@ -4140,19 +3868,17 @@ mod tests { sync::{Arc, Mutex as StdMutex}, }; - use oauth2::{AuthType, CsrfToken, HttpResponse, PkceCodeVerifier, TokenResponse}; + use oauth2::{AuthType, CsrfToken, HttpResponse, PkceCodeVerifier}; use reqwest::StatusCode; use rstest::rstest; - use tokio::sync::{Mutex, OwnedMutexGuard, Semaphore}; use url::Url; use super::{ AuthError, AuthorizationCallback, AuthorizationManager, AuthorizationMetadata, - AuthorizationMetadataSource, AuthorizationRequest, AuthorizationSession, - CredentialRefreshGuard, CredentialStore, InMemoryCredentialStore, InMemoryStateStore, - OAuthClientConfig, OAuthHttpClient, OAuthHttpClientError, OAuthHttpClientFuture, - OAuthHttpRedirectPolicy, OAuthHttpRequest, ScopeUpgradeConfig, StateStore, - StoredAuthorizationState, is_https_url, + AuthorizationMetadataSource, AuthorizationRequest, AuthorizationSession, CredentialStore, + InMemoryCredentialStore, InMemoryStateStore, OAuthClientConfig, OAuthHttpClient, + OAuthHttpClientError, OAuthHttpClientFuture, OAuthHttpRedirectPolicy, OAuthHttpRequest, + ScopeUpgradeConfig, StateStore, StoredAuthorizationState, is_https_url, }; use crate::transport::auth::VendorExtraTokenFields; @@ -4238,58 +3964,6 @@ mod tests { ); } - #[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(); @@ -4614,6 +4288,13 @@ mod tests { "authorization_servers": ["https://auth.example.com/tenant1"] }), ), + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/", + "authorization_servers": ["https://auth.example.com/tenant1"] + }), + ), http_response( 200, serde_json::json!({ @@ -4648,6 +4329,7 @@ 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", ], ) @@ -4665,6 +4347,13 @@ mod tests { "authorization_servers": ["https://auth.example.com/tenant1/"] }), ), + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/", + "authorization_servers": ["https://auth.example.com/tenant1/"] + }), + ), http_response( 200, serde_json::json!({ @@ -5128,579 +4817,6 @@ 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, @@ -8385,40 +7501,6 @@ 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; @@ -9190,349 +8272,4 @@ 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 deleted file mode 100644 index 226c8d5..0000000 --- a/crates/rmcp/src/transport/auth/enterprise.rs +++ /dev/null @@ -1,767 +0,0 @@ -//! 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 deleted file mode 100644 index 31a63c1..0000000 --- a/crates/rmcp/src/transport/auth/enterprise_tests.rs +++ /dev/null @@ -1,1145 +0,0 @@ -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/common/auth/streamable_http_client.rs b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs index 47b6caf..2053920 100644 --- a/crates/rmcp/src/transport/common/auth/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs @@ -19,10 +19,7 @@ where /// 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. + /// retry, then the challenge propagates. async fn call_reacting_to_challenges( &self, auth_token: Option, @@ -57,13 +54,10 @@ where 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(error) => { + debug!("token refresh after server rejection failed: {error}"); Err(StreamableHttpError::AuthRequired(challenge)) } - Err(error) => Err(error.into()), } } result => result, @@ -214,167 +208,3 @@ where .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 2ee37aa..e668d63 100644 --- a/crates/rmcp/src/transport/common/client_side_sse.rs +++ b/crates/rmcp/src/transport/common/client_side_sse.rs @@ -205,11 +205,6 @@ 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 { @@ -221,7 +216,6 @@ impl Default for ExponentialBackoff { Self { max_times: None, base_duration: Self::DEFAULT_DURATION, - max_delay: None, } } } @@ -233,16 +227,7 @@ impl SseRetryPolicy for ExponentialBackoff { { return None; } - // `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, - }) + Some(self.base_duration * (2u32.pow(current_times as u32))) } } @@ -790,75 +775,4 @@ 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 1b47560..7f08b6c 100644 --- a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs @@ -293,11 +293,6 @@ 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}"), ))); @@ -308,29 +303,17 @@ 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()) => { - // 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(_) - ) => - { + // 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) => { 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 688c9ad..5f995db 100644 --- a/crates/rmcp/src/transport/common/unix_socket.rs +++ b/crates/rmcp/src/transport/common/unix_socket.rs @@ -274,11 +274,6 @@ 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}"), ))); @@ -322,23 +317,13 @@ impl StreamableHttpClient for UnixSocketHttpClient { .map_err(|e| StreamableHttpError::Client(UnixSocketError::Hyper(e)))? .to_bytes(); 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(_) - ) => - { + Ok(message) => Ok(StreamableHttpPostResponse::Json(message, session_id)), + Err(e) => { 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 8af045f..fe563ee 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -1,16 +1,12 @@ use std::{ borrow::Cow, - collections::{HashMap, HashSet, VecDeque}, + collections::{HashMap, HashSet}, sync::Arc, time::Duration, }; -use futures::{ - Stream, StreamExt, - future::BoxFuture, - stream::{BoxStream, FuturesUnordered}, -}; -use http::{HeaderName, HeaderValue, StatusCode}; +use futures::{Stream, StreamExt, future::BoxFuture, stream::BoxStream}; +use http::{HeaderName, HeaderValue}; pub use sse_stream::Error as SseError; use sse_stream::Sse; use thiserror::Error; @@ -30,10 +26,7 @@ use crate::{ service::InboundStreamOrigin, transport::{ common::{client_side_sse::SseAutoReconnectStream, mcp_headers}, - worker::{ - RequestCancellationRegistration, Worker, WorkerQuitReason, WorkerSendRequest, - WorkerTransport, - }, + worker::{Worker, WorkerQuitReason, WorkerSendRequest, WorkerTransport}, }, }; @@ -214,41 +207,6 @@ pub enum StreamableHttpError { 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 { @@ -361,53 +319,12 @@ 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( @@ -558,21 +475,6 @@ 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 { @@ -592,85 +494,6 @@ 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()), @@ -686,16 +509,28 @@ 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, - ) -> Option { - let response_id = Self::server_response_id(message)?; - if let Some(id) = pending_stream_response_ids.take(response_id) { - return Some(id); + ) { + 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)); } - let id = RequestId::Number(response_id.numeric_string_value()?); - pending_stream_response_ids.take(&id) } async fn drain_queued_stream_messages( @@ -706,8 +541,7 @@ impl StreamableHttpClientWorker { loop { match sse_worker_rx.try_recv() { Ok(message) => { - let _ = - Self::clear_stream_response_pending(pending_stream_response_ids, &message); + 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(()), @@ -739,44 +573,6 @@ 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 @@ -835,34 +631,10 @@ 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, + + Send + + 'static, sse_worker_tx: tokio::sync::mpsc::Sender, origin: InboundStreamOrigin, close_on_response: bool, @@ -1047,19 +819,6 @@ 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 } @@ -1085,7 +844,6 @@ impl Worker for StreamableHttpClientWorker { let WorkerSendRequest { responder, message: startup_request, - .. } = context.recv_from_handler().await?; let is_legacy_startup = matches!( &startup_request, @@ -1195,31 +953,17 @@ 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>), - ControlMessage(WorkerSendRequest>), - StartPost(WorkerSendRequest>), - PostResult(PostResult), - RecoveryTimeout, + enum Event { + ClientMessage(WorkerSendRequest), 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 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 request_stream_cancellations = HashMap::::new(); let mut awaiting_fallback_initialized = false; if let Some(session_id) = &session_id { Self::spawn_common_stream( @@ -1232,159 +976,19 @@ impl Worker for StreamableHttpClientWorker { transport_task_ct.clone(), ); } - // Each POST uses the session and headers chosen when it starts. - // Only this loop updates the current session and protocol version. + // Main event loop - capture exit reason so we can do cleanup before returning 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.from_handler_rx.recv(), if may_receive => { + message = context.recv_from_handler() => { match message { - Some(msg) => Event::ClientMessage(msg), - None => break 'main_loop Err(WorkerQuitReason::HandlerTerminated), + Ok(msg) => Event::ClientMessage(msg), + Err(e) => break 'main_loop Err(e), } }, - 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"); @@ -1394,90 +998,48 @@ impl Worker for StreamableHttpClientWorker { }, terminated_stream = streams.join_next(), if !streams.is_empty() => { match terminated_stream { - Some(Ok((request_id, result))) => { - Event::StreamResult { request_id, result } + 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(Err(error)) => Event::StreamResult { - request_id: None, - result: Err(StreamableHttpError::TokioJoinError(error)), - }, - None => continue, } } }; match event { Event::ClientMessage(send_request) => { - 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 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, + } } - } - 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); + _ => 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(); } - let _ = send_request.responder.send(Ok(())); + pending_stream_response_ids.remove(&request_id); + let _ = 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, @@ -1486,9 +1048,6 @@ 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. @@ -1551,17 +1110,27 @@ impl Worker for StreamableHttpClientWorker { continue; } - let barrier = Self::is_ordering_barrier(message, &negotiated_version); - debug_assert!(!barrier || (posts.is_empty() && control_posts.is_empty())); - let inline_version = match message { + let request_id = Self::client_request_id(&message); + 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, ); @@ -1575,87 +1144,186 @@ impl Worker for StreamableHttpClientWorker { cleanup.protocol_headers = protocol_headers.clone(); } } - barrier_in_flight = barrier; - // The common stream can return a response before this POST finishes. - if let Some(request_id) = request_id { - pending_stream_response_ids.insert(request_id); - } - posts.push(Self::post_request( - self.client.clone(), - &config, - send_request, - PostSession { - id: session_id.clone(), - headers: request_headers, - version: request_version, - cancellation: session_cancellation.clone(), - }, - transport_task_ct.clone(), - )); - } - Event::PostResult(PostResult { - send_request, - response, - version, - }) => { - let is_control = Self::is_control_message(&send_request.message); - if !is_control { - // An ordering barrier runs without other ordinary POSTs. - barrier_in_flight = false; - } - let request_id = Self::client_request_id(&send_request.message); - if request_id - .as_ref() - .is_some_and(|id| !pending_stream_response_ids.contains(id)) - { - // A stream response or cancellation already completed this request. - let _ = send_request.responder.send(Ok(())); - continue; - } - let will_retry = - matches!(&response, Some(Err(StreamableHttpError::SessionExpired))) - && !is_control - && !retrying_recovery - && config.reinit_on_expired_session - && saved_init_request.is_some(); - let awaits_stream_response = matches!( - &response, - Some(Ok(StreamableHttpPostResponse::Accepted - | StreamableHttpPostResponse::Sse(..))) - ); - if !will_retry - && !awaits_stream_response - && let Some(id) = &request_id - { - pending_stream_response_ids.remove(id); - } - let Some(response) = response else { - let _ = send_request.responder.send(Ok(())); - continue; - }; - if will_retry { - if recovery_posts.is_empty() { - recovery_deadline = - Some(tokio::time::Instant::now() + config.session_recovery_timeout); - } - recovery_posts.push_back(send_request); - continue; - } - let request_cancellation = send_request.cancellation_registration(); - let WorkerSendRequest { - message, responder, .. - } = send_request; - let is_initialized_notification = matches!( - &message, - ClientJsonRpcMessage::Notification(notification) - if matches!( - ¬ification.notification, - ClientNotification::InitializedNotification(_) - ) - ); + let 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(); + 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, origin, 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) + } + } 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(()) } @@ -1663,13 +1331,17 @@ impl Worker for StreamableHttpClientWorker { cache_tools_from_response( &mut tool_header_cache, &mut message, - &version, + &negotiated_version, ); context.send_to_handler(message).await?; Ok(()) } Ok(StreamableHttpPostResponse::Sse(stream, ..)) => { - let stream_request_id = request_id; + 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(), @@ -1680,24 +1352,11 @@ impl Worker for StreamableHttpClientWorker { config.max_sse_event_size, self.config.retry_config.clone(), ); - let request_ct = request_cancellation - .as_ref() - .map(|registration| registration.token()) - .unwrap_or_else(|| transport_task_ct.child_token()); - // A legacy adapter may need the open stream to handle cancellation. - let stream_ct = if uses_modern_http { - request_ct.clone() - } else { - request_cancellation - .as_ref() - .map(|registration| registration.lifetime_token()) - .unwrap_or_else(|| request_ct.clone()) - }; - if let (Some(request_id), Some(registration)) = - (stream_request_id.as_ref(), request_cancellation) + 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(), registration); + .insert(request_id.clone(), stream_ct.clone()); } let stream_tx = sse_worker_tx.clone(); let origin = match &stream_request_id { @@ -1705,13 +1364,8 @@ impl Worker for StreamableHttpClientWorker { None => InboundStreamOrigin::Unassociated, }; streams.spawn(async move { - let result = Self::run_response_stream( - sse_stream, - stream_tx, - origin, - request_ct, - stream_ct, - uses_modern_http, + let result = Self::execute_sse_stream( + sse_stream, stream_tx, origin, true, stream_ct, ) .await; (stream_request_id, result) @@ -1740,13 +1394,18 @@ impl Worker for StreamableHttpClientWorker { let _ = responder.send(send_result); } Event::ServerMessage(mut json_rpc_message) => { - // Match against all pending requests, not just open response streams. - if let Some(request_id) = Self::clear_stream_response_pending( + 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( &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, @@ -1765,10 +1424,8 @@ impl Worker for StreamableHttpClientWorker { &mut pending_stream_response_ids, ) .await?; - let cancelled = request_stream_cancellations - .remove(&request_id) - .is_some_and(|registration| registration.token().is_cancelled()); - if pending_stream_response_ids.remove(&request_id) && !cancelled { + request_stream_cancellations.remove(&request_id); + if pending_stream_response_ids.remove(&request_id) { context .send_to_handler(ServerJsonRpcMessage::error( ErrorData::transport_closed( @@ -1789,14 +1446,6 @@ 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 { @@ -2029,19 +1678,6 @@ 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 @@ -2054,17 +1690,15 @@ pub struct StreamableHttpClientTransportConfig { /// [`StreamableHttpClient`] implementations must override the corresponding /// `*_with_max_sse_event_size` methods to enforce it. pub max_sse_event_size: usize, - /// Automatically creates a new session when the server reports an expired - /// session (`http 404`). + /// Enables transparent recovery when the server reports an expired session (`HTTP 404`). /// - /// 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. + /// 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`. /// - /// Control POSTs and other POST failures are not retried. If recovery or a retry - /// fails, the transport returns an error to the caller. + /// 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. pub reinit_on_expired_session: bool, } @@ -2076,24 +1710,6 @@ 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 @@ -2158,9 +1774,6 @@ 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(), @@ -2184,32 +1797,6 @@ mod tests { 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" @@ -2348,7 +1935,6 @@ mod tests { Arc::new(ExponentialBackoff { max_times: Some(1), base_duration: Duration::ZERO, - max_delay: None, }), ); let mut stream = std::pin::pin!(stream); @@ -2453,7 +2039,6 @@ mod tests { Arc::new(ExponentialBackoff { max_times: Some(1), base_duration: Duration::ZERO, - max_delay: None, }), ); @@ -2573,13 +2158,11 @@ mod tests { NumberOrString::String("1".into()), ); - let matched_id = - StreamableHttpClientWorker::::clear_stream_response_pending( - &mut pending, - &response, - ); + StreamableHttpClientWorker::::clear_stream_response_pending( + &mut pending, + &response, + ); - assert_eq!(matched_id, Some(NumberOrString::Number(1))); assert!(pending.is_empty()); } @@ -2590,16 +2173,14 @@ mod tests { let mut pending = HashSet::from([NumberOrString::Number(1), string_id.clone()]); let response = ServerJsonRpcMessage::response( ServerResult::ListToolsResult(ListToolsResult::default()), - string_id.clone(), + string_id, ); - let matched_id = - StreamableHttpClientWorker::::clear_stream_response_pending( - &mut pending, - &response, - ); + 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 e03e5b7..cc9e148 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/local.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/local.rs @@ -1127,9 +1127,7 @@ 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 1b53820..f1fef58 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/crates/rmcp/src/transport/streamable_http_server/tower.rs @@ -28,7 +28,7 @@ use crate::{ ClientCapabilities, ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorCode, ErrorData, GetExtensions, GetMeta, Implementation, InitializeRequest, InitializeRequestParams, InitializedNotification, JsonObject, JsonRpcError, - ProtocolVersion, RequestId, ServerConfig, ServerJsonRpcMessage, ServerResult, + ProtocolVersion, RequestId, ServerInfo, ServerJsonRpcMessage, ServerResult, }, serve_server, service::{ @@ -55,24 +55,6 @@ 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 { @@ -110,18 +92,13 @@ pub struct StreamableHttpServerConfig { pub allowed_hosts: Vec, /// Allowed browser origins for inbound `Origin` validation. /// - /// 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. + /// 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`. /// 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 @@ -194,7 +171,6 @@ 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, @@ -222,15 +198,9 @@ impl StreamableHttpServerConfig { self.allowed_origins = allowed_origins.into_iter().map(Into::into).collect(); self } - /// 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. + /// Disable Origin validation, reverting to the default ignore-Origin behavior. 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 { @@ -277,6 +247,10 @@ impl StreamableHttpServerConfig { } } +#[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: @@ -285,7 +259,7 @@ impl StreamableHttpServerConfig { fn validate_protocol_version_header( headers: &http::HeaderMap, allow_unknown: bool, -) -> HttpResult<()> { +) -> Result<(), BoxResponse> { if let Some(value) = headers.get(HEADER_MCP_PROTOCOL_VERSION) { let version_str = value.to_str().map_err(|_| { Response::builder() @@ -310,8 +284,7 @@ fn validate_protocol_version_header( ))) .boxed(), ) - .expect("valid response") - .into()); + .expect("valid response")); } } Ok(()) @@ -349,7 +322,7 @@ impl> Service for NegotiatingStatelessHttpSer &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(); @@ -367,7 +340,7 @@ impl> Service for NegotiatingStatelessHttpSer self.0.handle_notification(notification, context).await } - fn get_info(&self) -> ServerConfig { + fn get_info(&self) -> ServerInfo { self.0.get_info() } @@ -376,52 +349,48 @@ impl> Service for NegotiatingStatelessHttpSer } } +#[expect( + clippy::result_large_err, + reason = "BoxResponse is intentionally large; matches other handlers in this file" +)] // SEP-2567: sessions are removed from the discover lifecycle. Validate // protocol-version consistency, then classify the request with the shared // lifecycle helper. fn is_legacy_request( message: Option<&ClientJsonRpcMessage>, headers: &HeaderMap, -) -> HttpResult { +) -> Result { let has_per_request_version = message.is_some_and(message_has_per_request_protocol_version); validate_protocol_version_header(headers, has_per_request_version)?; if let Some(message) = message { - if let ClientJsonRpcMessage::Request(req) = message - && let ClientRequest::InitializeRequest(init) = &req.request - { - validate_header_matches_init_body( - headers, - init.params.protocol_version.as_str(), - Some(req.id.clone()), - )?; + 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()), + )?; + } } validate_request_protocol_version_meta(headers, message)?; } - // An `initialize` request selects legacy semantics whatever version it names: - // the handshake exists only in the revisions before 2026-07-28, so the - // version in its params never routes it to the stateless path. The - // handshake itself answers with a legacy version the server supports. - if matches!( - message, - Some(ClientJsonRpcMessage::Request(req)) - if matches!(&req.request, ClientRequest::InitializeRequest(_)) - ) { - return Ok(true); - } - let uses_discover_lifecycle = matches!( message, Some(ClientJsonRpcMessage::Request(req)) - if req - .request - .get_meta() - .missing_required_keys(&ProtocolVersion::V_2026_07_28) - .is_empty() + if !matches!(&req.request, ClientRequest::InitializeRequest(_)) + && req + .request + .get_meta() + .missing_required_keys(&ProtocolVersion::V_2026_07_28) + .is_empty() ); let from_body = match message { - Some(ClientJsonRpcMessage::Request(req)) => req.request.get_meta().protocol_version(), + Some(ClientJsonRpcMessage::Request(req)) => match &req.request { + ClientRequest::InitializeRequest(init) => Some(init.params.protocol_version.clone()), + _ => req.request.get_meta().protocol_version(), + }, _ => None, }; let version = from_body @@ -453,10 +422,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 - && sender.send(event).await.is_err() - { - *output = None; + if let Some(sender) = output { + if sender.send(event).await.is_err() { + *output = None; + } } Ok(()) } @@ -487,12 +456,16 @@ 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, -) -> HttpResult<()> { +) -> Result<(), BoxResponse> { let Some(header_value) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else { return Ok(()); }; @@ -513,16 +486,19 @@ 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, -) -> HttpResult<()> { +) -> Result<(), BoxResponse> { let ClientJsonRpcMessage::Request(request) = message else { return Ok(()); }; @@ -546,8 +522,7 @@ fn validate_request_protocol_version_meta( "Invalid params: request _meta is missing or has malformed required fields: {}", missing.join(", ") ), - ) - .into()); + )); } return Ok(()); }; @@ -555,8 +530,7 @@ fn validate_request_protocol_version_meta( return Err(header_mismatch_jsonrpc_response( Some(request.id.clone()), "request _meta protocolVersion requires MCP-Protocol-Version header", - ) - .into()); + )); }; if header_version != meta_version.as_str() { return Err(header_mismatch_jsonrpc_response( @@ -564,8 +538,7 @@ fn validate_request_protocol_version_meta( format!( "MCP-Protocol-Version header ({header_version}) does not match request _meta protocolVersion ({meta_version})" ), - ) - .into()); + )); } Ok(()) } @@ -576,11 +549,15 @@ fn validate_request_protocol_version_meta( /// 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. +#[expect( + clippy::result_large_err, + reason = "BoxResponse is intentionally large; matches other handlers in this file" +)] fn validate_required_protocol_header( config: &StreamableHttpServerConfig, headers: &HeaderMap, message: &ClientJsonRpcMessage, -) -> HttpResult<()> { +) -> Result<(), BoxResponse> { if !config.stateless_protocol_metadata_required { return Ok(()); } @@ -598,8 +575,7 @@ fn validate_required_protocol_header( 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, @@ -609,10 +585,14 @@ fn validate_required_protocol_header( /// `server/discover` (whose body-metadata rule is already enforced by /// `validate_request_protocol_version_meta`), notifications, and other message /// kinds are exempt. +#[expect( + clippy::result_large_err, + reason = "BoxResponse is intentionally large; matches other handlers in this file" +)] fn validate_required_protocol_meta( config: &StreamableHttpServerConfig, message: &ClientJsonRpcMessage, -) -> HttpResult<()> { +) -> Result<(), BoxResponse> { if !config.stateless_protocol_metadata_required { return Ok(()); } @@ -631,8 +611,7 @@ fn validate_required_protocol_meta( 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 { @@ -644,56 +623,16 @@ fn jsonrpc_http_status(message: &ServerJsonRpcMessage) -> http::StatusCode { match error.error.code { ErrorCode::UNSUPPORTED_PROTOCOL_VERSION | ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY - | ErrorCode::INVALID_PARAMS - | ErrorCode::HEADER_MISMATCH => http::StatusCode::BAD_REQUEST, + | ErrorCode::INVALID_PARAMS => http::StatusCode::BAD_REQUEST, ErrorCode::METHOD_NOT_FOUND => http::StatusCode::NOT_FOUND, _ => http::StatusCode::OK, } } -#[cfg(test)] -mod jsonrpc_http_status_tests { - use super::*; - - fn error_message(code: ErrorCode) -> ServerJsonRpcMessage { - ServerJsonRpcMessage::Error(JsonRpcError::new( - Some(RequestId::Number(1)), - ErrorData::new(code, "test error", None), - )) - } - - /// A handler returning `ErrorData::header_mismatch(..)` on the modern per-request path - /// must map to HTTP 400, not the default HTTP 200. Regression test for - /// https://github.com/modelcontextprotocol/rust-sdk/issues/1225 - #[test] - fn header_mismatch_maps_to_bad_request() { - assert_eq!( - jsonrpc_http_status(&error_message(ErrorCode::HEADER_MISMATCH)), - http::StatusCode::BAD_REQUEST - ); - } - - #[test] - fn method_not_found_maps_to_not_found() { - assert_eq!( - jsonrpc_http_status(&error_message(ErrorCode::METHOD_NOT_FOUND)), - http::StatusCode::NOT_FOUND - ); - } - - #[test] - fn unmapped_error_defaults_to_ok() { - assert_eq!( - jsonrpc_http_status(&error_message(ErrorCode::INTERNAL_ERROR)), - http::StatusCode::OK - ); - } -} - fn jsonrpc_message_response( message: ServerJsonRpcMessage, map_protocol_status: bool, -) -> HttpResult { +) -> Result { let status = if map_protocol_status { jsonrpc_http_status(&message) } else { @@ -727,11 +666,15 @@ 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>, -) -> HttpResult<()> { +) -> Result<(), BoxResponse> { let version_requires_headers = headers .get(HEADER_MCP_PROTOCOL_VERSION) .and_then(|value| value.to_str().ok()) @@ -764,7 +707,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).into()); + return Err(header_mismatch_jsonrpc_response(request_id, reason)); } Ok(()) } @@ -854,6 +797,9 @@ 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)) @@ -885,7 +831,10 @@ fn bad_request_response(message: &str) -> BoxResponse { .expect("failed to build bad request response") } -fn parse_host_header(uri: &http::Uri, headers: &HeaderMap) -> HttpResult { +fn parse_host_header( + uri: &http::Uri, + headers: &HeaderMap, +) -> Result { if let Some(host) = headers.get(http::header::HOST) { let host_str = host .to_str() @@ -916,24 +865,24 @@ fn validate_dns_rebinding_headers( uri: &http::Uri, headers: &HeaderMap, config: &StreamableHttpServerConfig, -) -> HttpResult<()> { +) -> Result<(), BoxResponse> { 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").into()); + return Err(forbidden_response("Forbidden: Host header is not allowed")); } - validate_origin_header(headers, config)?; + validate_origin_header(headers, &config.allowed_origins)?; Ok(()) } fn validate_origin_header( headers: &HeaderMap, - config: &StreamableHttpServerConfig, -) -> HttpResult<()> { - if !config.validate_empty_origin_allowlist && config.allowed_origins.is_empty() { + allowed_origins: &[String], +) -> Result<(), BoxResponse> { + if allowed_origins.is_empty() { return Ok(()); } let Some(origin_header) = headers.get(http::header::ORIGIN) else { @@ -944,20 +893,22 @@ fn validate_origin_header( .inspect_err(|_| { tracing::warn!(origin = ?origin_header, "rejected request with non-UTF-8 Origin header"); }) - .map_err(|_| forbidden_response("Forbidden: Invalid Origin header encoding"))?; + .map_err(|_| bad_request_response("Bad Request: 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", ); - forbidden_response("Forbidden: Invalid Origin header") + bad_request_response("Bad Request: Invalid Origin header") })?; - if !origin_is_allowed(&origin, &config.allowed_origins) { + if !origin_is_allowed(&origin, 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").into()); + return Err(forbidden_response( + "Forbidden: Origin header is not allowed", + )); } Ok(()) } @@ -1053,7 +1004,9 @@ 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, + pending_restores: Option< + Arc>>>>, + >, /// 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. @@ -1106,9 +1059,10 @@ where /// `result` defaults to `false` (failure / cancellation). Only the success path /// needs to set it to `true` before returning. struct PendingRestoreGuard { - pending_restores: PendingRestores, + pending_restores: + Arc>>>>, session_id: SessionId, - watch_tx: RestoreResultSender, + watch_tx: tokio::sync::watch::Sender>, /// The value that will be broadcast to waiting tasks on drop. result: bool, } @@ -1136,10 +1090,12 @@ where session_manager: Arc, config: StreamableHttpServerConfig, ) -> Self { - let pending_restores = config - .session_store - .is_some() - .then(|| Arc::new(tokio::sync::RwLock::new(HashMap::new()))); + let pending_restores = config.session_store.is_some().then(|| { + Arc::new(tokio::sync::RwLock::new(HashMap::< + SessionId, + tokio::sync::watch::Sender>, + >::new())) + }); Self { config, session_manager, @@ -1166,18 +1122,19 @@ where tokio::spawn(async move { let mut sender = Some(sender); - if let Some(retry) = retry - && let Err(error) = persist_and_forward_event( + if let Some(retry) = retry { + if 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; @@ -1249,7 +1206,7 @@ where service: S, mut request: crate::model::JsonRpcRequest, parts: http::request::Parts, - ) -> HttpResult { + ) -> Result { let peer_info = Self::peer_info_for_stateless_request(&request, &parts.headers); request.request.extensions_mut().insert(parts); let (transport, mut receiver) = @@ -1314,10 +1271,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() - && let Some(schema) = cache.get(name) - { - return schema.clone(); + if let Ok(cache) = self.tool_schemas.read() { + if let Some(schema) = cache.get(name) { + return schema.clone(); + } } let schema = self .get_service() @@ -1507,15 +1464,23 @@ where Some(init_done_tx), ); - self.session_manager + if let Err(e) = self + .session_manager .initialize_session(session_id, restore_init) .await - .map_err(|e| std::io::Error::other(e.to_string()))?; + .map_err(|e| std::io::Error::other(e.to_string())) + { + return Err(e); + } - self.session_manager + if let Err(e) = self + .session_manager .accept_message(session_id, restore_initialized) .await - .map_err(|e| std::io::Error::other(e.to_string()))?; + .map_err(|e| std::io::Error::other(e.to_string())) + { + return Err(e); + } if init_done_rx.await.is_err() { return Err(std::io::Error::other( @@ -1540,7 +1505,7 @@ where if let Err(response) = validate_dns_rebinding_headers(request.uri(), request.headers(), &self.config) { - return response.into_response(); + return response; } let method = request.method().clone(); let supports_stateless_replay = self.session_manager.event_store().is_some(); @@ -1567,10 +1532,10 @@ where }; match result { Ok(response) => response, - Err(response) => response.into_response(), + Err(response) => response, } } - async fn handle_get(&self, request: Request) -> HttpResult + async fn handle_get(&self, request: Request) -> Result where B: Body + Send + 'static, B::Error: Display, @@ -1713,7 +1678,7 @@ where )) } - async fn handle_post(&self, request: Request) -> HttpResult + async fn handle_post(&self, request: Request) -> Result where B: Body + Send + 'static, B::Error: Display, @@ -1865,7 +1830,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").into()); + return Err(unexpected_message_response("initialize request")); }; // Reject mismatched MCP-Protocol-Version header before binding the session to anything. validate_header_matches_init_body( @@ -1883,7 +1848,7 @@ where stored_init_params } _ => { - return Err(unexpected_message_response("initialize request").into()); + return Err(unexpected_message_response("initialize request")); } }; let service = self @@ -2040,8 +2005,7 @@ where std::io::ErrorKind::UnexpectedEof, "no response message received from handler", ), - ) - .into()); + )); }; tracing::trace!(?message); if matches!( @@ -2073,7 +2037,7 @@ where } } - async fn handle_delete(&self, request: Request) -> HttpResult + async fn handle_delete(&self, request: Request) -> Result where B: Body + Send + 'static, B::Error: Display, @@ -2111,7 +2075,7 @@ where Ok(accepted_response()) } - /// Build a `InitializeRequestParams` (peer_info) for a stateless request so that + /// Build a `ClientInfo` (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 e32da7d..5294640 100644 --- a/crates/rmcp/src/transport/worker.rs +++ b/crates/rmcp/src/transport/worker.rs @@ -1,21 +1,10 @@ -use std::{ - borrow::Cow, - collections::HashMap, - sync::{ - Arc, Mutex, PoisonError, Weak, - atomic::{AtomicU64, Ordering}, - }, - time::Duration, -}; +use std::{borrow::Cow, time::Duration}; use tokio_util::sync::CancellationToken; use tracing::{Instrument, Level}; use super::{IntoTransport, Transport}; -use crate::{ - model::{CancelledNotification, JsonRpcMessage, RequestId}, - service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage}, -}; +use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage}; #[derive(Debug, thiserror::Error)] #[non_exhaustive] @@ -64,114 +53,17 @@ 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, @@ -212,17 +104,11 @@ 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(), }; @@ -256,32 +142,11 @@ 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] @@ -294,32 +159,10 @@ 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, @@ -347,52 +190,15 @@ impl Transport for WorkerTransport { &mut self, item: TxJsonRpcMessage, ) -> impl Future> + Send + 'static { - 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 tx = self.send_service.clone(); 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(()) } } @@ -409,133 +215,3 @@ 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 961d977..22c6d38 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, ServerConfig}, + model::{ServerCapabilities, ServerInfo}, schemars, tool, tool_router, }; #[derive(Debug, serde::Deserialize, schemars::JsonSchema)] @@ -52,8 +52,8 @@ impl Calculator { } impl ServerHandler for Calculator { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 7e68621..7928fa6 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) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_logging().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_logging().build()) } #[allow(deprecated)] diff --git a/crates/rmcp/tests/test_cancelled_response.rs b/crates/rmcp/tests/test_cancelled_response.rs index 864932a..5961a7b 100644 --- a/crates/rmcp/tests/test_cancelled_response.rs +++ b/crates/rmcp/tests/test_cancelled_response.rs @@ -3,27 +3,13 @@ //! until the request is cancelled, so its result is only produced *after* the //! cancellation — the service loop must drop it rather than write it to the wire. -#[cfg(all(feature = "client", not(feature = "local")))] -use std::sync::Arc; use std::{collections::BTreeSet, process::Stdio, time::Duration}; -#[cfg(all(feature = "client", not(feature = "local")))] -use rmcp::{ - ClientHandler, RoleClient, - model::{ - CancelledNotification, CancelledNotificationParam, ClientJsonRpcMessage, ClientRequest, - ClientResult, ElicitRequest, ElicitRequestParams, ElicitResult, ElicitationAction, - ElicitationSchema, PingRequest, RequestId, ServerJsonRpcMessage, ServerNotification, - ServerRequest, ServerResult, - }, - service::{PeerRequestOptions, QuitReason, serve_directly}, - transport::{IntoTransport, Transport}, -}; use rmcp::{ ErrorData as McpError, RoleServer, ServerHandler, ServiceExt, model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ServerCapabilities, - ServerConfig, + ServerInfo, }, service::RequestContext, }; @@ -35,7 +21,6 @@ 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<()> { @@ -58,7 +43,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" }), @@ -71,7 +56,7 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> { &mut writer, &json!({ "jsonrpc": "2.0", - "id": STRING_REQUEST_ID, + "id": 2, "method": "tools/call", "params": { "name": "wait-for-cancel", "arguments": {} } }), @@ -82,59 +67,21 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> { &json!({ "jsonrpc": "2.0", "method": "notifications/cancelled", - "params": { "requestId": "unrelated-request" } - }), - ) - .await?; - send_json( - &mut writer, - &json!({ - "jsonrpc": "2.0", - "method": "notifications/cancelled", - "params": { "requestId": STRING_REQUEST_ID } + "params": { "requestId": 2 } }), ) .await?; // A ping proves the server is alive past the cancellation, so the absence of - // a cancelled response is genuine suppression rather than a dead connection. + // an id=2 response is genuine suppression rather than a dead connection. send_json( &mut writer, &json!({ "jsonrpc": "2.0", "id": 3, "method": "ping" }), ) .await?; - let seen = collect_ids_until(&mut reader, "3", READ_TIMEOUT).await?; - assert!(seen.contains("3")); - assert!(!seen.contains(STRING_REQUEST_ID)); - - send_json( - &mut writer, - &json!({ - "jsonrpc": "2.0", - "id": 4, - "method": "tools/call", - "params": { "name": "complete-unless-cancelled", "arguments": {} } - }), - ) - .await?; - send_json( - &mut writer, - &json!({ - "jsonrpc": "2.0", - "method": "notifications/cancelled", - "params": { "requestId": "4" } - }), - ) - .await?; - send_json( - &mut writer, - &json!({ "jsonrpc": "2.0", "id": 5, "method": "ping" }), - ) - .await?; - - let seen = collect_ids_until(&mut reader, "5", READ_TIMEOUT).await?; - assert!(seen.contains("4"), "string and numeric IDs must not alias"); - assert!(seen.contains("5")); + let seen = collect_ids_until(&mut reader, 3, READ_TIMEOUT).await?; + assert!(seen.contains(&3)); + assert!(!seen.contains(&2)); drop(writer); wait_for_child(&mut child).await; @@ -144,383 +91,20 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> { struct WaitForCancelServer; impl ServerHandler for WaitForCancelServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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") { @@ -588,9 +172,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: &str, + stop_id: u64, timeout: Duration, -) -> anyhow::Result> +) -> anyhow::Result> where R: tokio::io::AsyncRead + Unpin, { @@ -616,13 +200,8 @@ where let Ok(value) = serde_json::from_str::(trimmed) else { continue; }; - let id = match value.get("id") { - Some(Value::String(id)) => Some(id.clone()), - Some(Value::Number(id)) => Some(id.to_string()), - _ => None, - }; - if let Some(id) = id { - seen.insert(id.clone()); + if let Some(id) = value.get("id").and_then(Value::as_u64) { + seen.insert(id); 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_custom_headers.rs b/crates/rmcp/tests/test_custom_headers.rs index b01223a..736dce1 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-") - && let Ok(v) = value.to_str() - { - headers_map.insert(name_str.to_string(), v.to_string()); + if name_str.starts_with("x-") { + if 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) - && 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) { + 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" + } } - } - }); - 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, ServerConfig}, + model::{ServerCapabilities, ServerInfo}, 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) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().build()) } } @@ -880,11 +880,10 @@ fn test_protocol_version_utilities() { assert!(ProtocolVersion::KNOWN_VERSIONS.contains(&ProtocolVersion::V_2026_07_28)); } -/// Integration test: Verify Host validation remains enabled when Origin validation is -/// disabled by default +/// Integration test: Verify server validates only the Host header for DNS rebinding protection #[tokio::test] #[cfg(all(feature = "transport-streamable-http-server", feature = "server",))] -async fn test_server_validates_host_when_origin_validation_is_disabled_by_default() { +async fn test_server_validates_host_header_for_dns_rebinding_protection() { use std::sync::Arc; use bytes::Bytes; @@ -892,7 +891,7 @@ async fn test_server_validates_host_when_origin_validation_is_disabled_by_defaul use http_body_util::Full; use rmcp::{ handler::server::ServerHandler, - model::{ServerCapabilities, ServerConfig}, + model::{ServerCapabilities, ServerInfo}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -903,8 +902,8 @@ async fn test_server_validates_host_when_origin_validation_is_disabled_by_defaul struct TestHandler; impl ServerHandler for TestHandler { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().build()) } } @@ -975,7 +974,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, ServerConfig}, + model::{ServerCapabilities, ServerInfo}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -986,8 +985,8 @@ async fn test_server_validates_host_header_port_for_dns_rebinding_protection() { struct TestHandler; impl ServerHandler for TestHandler { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().build()) } } @@ -1046,7 +1045,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, ServerConfig}, + model::{ServerCapabilities, ServerInfo}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -1057,8 +1056,8 @@ async fn test_server_falls_back_to_uri_authority_when_host_header_missing() { struct TestHandler; impl ServerHandler for TestHandler { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().build()) } } @@ -1128,11 +1127,11 @@ mod origin_validation { use std::sync::Arc; use bytes::Bytes; - use http::{HeaderValue, Method, Request, header::CONTENT_TYPE}; + use http::{Method, Request, header::CONTENT_TYPE}; use http_body_util::Full; use rmcp::{ handler::server::ServerHandler, - model::{ServerCapabilities, ServerConfig}, + model::{ServerCapabilities, ServerInfo}, transport::streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }, @@ -1143,25 +1142,17 @@ mod origin_validation { struct TestHandler; impl ServerHandler for TestHandler { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 { - service_with_config( + StreamableHttpService::new( + || Ok(TestHandler), + Arc::new(LocalSessionManager::default()), StreamableHttpServerConfig::default().with_allowed_origins(origins.iter().copied()), ) } @@ -1208,80 +1199,6 @@ 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_discover_http_client_startup.rs b/crates/rmcp/tests/test_discover_http_client_startup.rs index 2d92df5..c6051a0 100644 --- a/crates/rmcp/tests/test_discover_http_client_startup.rs +++ b/crates/rmcp/tests/test_discover_http_client_startup.rs @@ -5,18 +5,11 @@ feature = "transport-streamable-http-server" ))] -use std::{borrow::Cow, sync::Arc}; +use std::borrow::Cow; -use axum::{ - Router, - body::{Body, Bytes}, - extract::State, - http::{Response, StatusCode}, - routing::post, -}; use rmcp::{ ClientLifecycleMode, ClientServiceExt, ServerHandler, - model::{ClientConfig, DiscoverResult, ErrorCode, ErrorData, ProtocolVersion}, + model::{ClientInfo, DiscoverResult, ErrorCode, ErrorData, ProtocolVersion}, service::{MaybeSendFuture, RequestContext, RoleServer}, transport::{ StreamableHttpClientTransport, @@ -26,8 +19,6 @@ use rmcp::{ }, }, }; -use serde_json::json; -use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; #[derive(Clone, Default)] @@ -55,54 +46,6 @@ 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(); @@ -132,7 +75,7 @@ async fn discover_http_client_bootstraps_headers_without_initialize() { let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(format!("http://{address}/mcp")), ); - let client = ClientConfig::default() + let client = ClientInfo::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { @@ -176,7 +119,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 = ClientConfig::default() + let client = ClientInfo::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Auto { @@ -192,47 +135,3 @@ 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_handler_cache_hints.rs b/crates/rmcp/tests/test_handler_cache_hints.rs index ed3cf67..d6b626d 100644 --- a/crates/rmcp/tests/test_handler_cache_hints.rs +++ b/crates/rmcp/tests/test_handler_cache_hints.rs @@ -2,15 +2,10 @@ #![cfg(feature = "client")] use rmcp::{ - ClientHandler, RoleClient, RoleServer, ServerHandler, + ClientHandler, ServerHandler, ServiceExt, handler::server::router::{prompt::PromptRouter, tool::ToolRouter}, - model::{ - CacheScope, ClientConfig, InitializeResult, ListPromptsResult, ListToolsResult, - ProtocolVersion, - }, - prompt_handler, - service::serve_directly, - tool_handler, + model::{CacheScope, ClientInfo, ListPromptsResult, ListToolsResult, ProtocolVersion}, + prompt_handler, tool_handler, }; #[derive(Debug, Clone)] @@ -38,40 +33,29 @@ struct VersionedClient { } impl ClientHandler for VersionedClient { - fn get_info(&self) -> ClientConfig { - let mut info = ClientConfig::default(); + fn get_info(&self) -> ClientInfo { + let mut info = ClientInfo::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?; + CacheHintServer::new() + .serve(server_transport) + .await? + .waiting() + .await?; anyhow::Ok(()) }); - let client = serve_directly::( - client_handler, - client_transport, - Some(server_peer_info.into()), - ); + let client = VersionedClient { protocol_version } + .serve(client_transport) + .await + .expect("client should connect"); let tools = client .list_tools(None) .await diff --git a/crates/rmcp/tests/test_inflight_response_drain.rs b/crates/rmcp/tests/test_inflight_response_drain.rs index fc5908f..c75e17c 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, ClientConfig, ServerCapabilities, ServerConfig}, + model::{CallToolRequestParams, ClientInfo, ServerCapabilities, ServerInfo}, service::QuitReason, tool, tool_handler, tool_router, }; @@ -55,8 +55,8 @@ impl SlowToolServer { #[tool_handler] impl ServerHandler for SlowToolServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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) -> ClientConfig { - ClientConfig::default() + fn get_info(&self) -> ClientInfo { + ClientInfo::default() } } diff --git a/crates/rmcp/tests/test_live_oauth_refresh.rs b/crates/rmcp/tests/test_live_oauth_refresh.rs deleted file mode 100644 index 543abe5..0000000 --- a/crates/rmcp/tests/test_live_oauth_refresh.rs +++ /dev/null @@ -1,233 +0,0 @@ -//! 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_message_schema/server_json_rpc_message_schema.json b/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json index e220078..ca0ef56 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 @@ -963,10 +963,6 @@ "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": [ diff --git a/crates/rmcp/tests/test_mrtr_behavior.rs b/crates/rmcp/tests/test_mrtr_behavior.rs index d9df5af..4fb24e1 100644 --- a/crates/rmcp/tests/test_mrtr_behavior.rs +++ b/crates/rmcp/tests/test_mrtr_behavior.rs @@ -114,8 +114,8 @@ impl MacroMrtrServer { #[tool_handler] impl ServerHandler for MacroMrtrServer { - fn get_info(&self) -> ServerConfig { - let mut info = ServerConfig::new(ServerCapabilities::builder().enable_tools().build()); + fn get_info(&self) -> ServerInfo { + let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build()); info.protocol_version = ProtocolVersion::V_2026_07_28; info } @@ -221,8 +221,8 @@ impl MrtrServer { } impl ServerHandler for MrtrServer { - fn get_info(&self) -> ServerConfig { - let mut info = ServerConfig::new( + fn get_info(&self) -> ServerInfo { + let mut info = ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_prompts() @@ -328,16 +328,16 @@ impl ClientHandler for MrtrClient { // Harness // ============================================================================= -fn client_info(protocol_version: ProtocolVersion) -> InitializeRequestParams { - InitializeRequestParams::new( +fn client_info(protocol_version: ProtocolVersion) -> ClientInfo { + ClientInfo::new( ClientCapabilities::builder().enable_elicitation().build(), Implementation::new("mrtr-test-client", "0.0.0"), ) .with_protocol_version(protocol_version) } -fn server_info(protocol_version: ProtocolVersion) -> InitializeResult { - let mut info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build()); +fn server_info(protocol_version: ProtocolVersion) -> ServerInfo { + let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build()); info.protocol_version = protocol_version; info } @@ -594,23 +594,19 @@ 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. - 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", - ); + static KEY: &[u8] = b"integration-signing-key-32-bytes!"; fn codec() -> &'static RequestStateCodec { static CODEC: OnceLock = OnceLock::new(); - CODEC.get_or_init(|| RequestStateCodec::new_unchecked(KEY)) + CODEC.get_or_init(|| RequestStateCodec::new(KEY)) } #[derive(Clone, Default)] struct SealingServer; impl ServerHandler for SealingServer { - fn get_info(&self) -> ServerConfig { - let mut info = ServerConfig::new(ServerCapabilities::builder().enable_tools().build()); + fn get_info(&self) -> ServerInfo { + let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build()); info.protocol_version = ProtocolVersion::V_2026_07_28; info } diff --git a/crates/rmcp/tests/test_notification.rs b/crates/rmcp/tests/test_notification.rs index dafb0df..9aafd3c 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, ServerConfig, ServerNotification, SubscribeRequestParams, + ServerCapabilities, ServerInfo, 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) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::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 e6d65b6..7a00249 100644 --- a/crates/rmcp/tests/test_prompt_macros.rs +++ b/crates/rmcp/tests/test_prompt_macros.rs @@ -4,12 +4,14 @@ use std::sync::Arc; use rmcp::{ - ClientHandler, ServerHandler, ServiceExt, + ClientHandler, RoleServer, ServerHandler, ServiceExt, handler::server::{router::prompt::PromptRouter, wrapper::Parameters}, model::{ - ClientConfig, ContentBlock, GetPromptRequestParams, GetPromptResult, PromptMessage, Role, + ClientInfo, ContentBlock, GetPromptRequestParams, GetPromptResult, ListPromptsResult, + PaginatedRequestParams, PromptMessage, Role, }, prompt, prompt_handler, prompt_router, + service::RequestContext, }; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; @@ -297,8 +299,8 @@ fn test_optional_field_schema_generation_via_macro() { struct DummyClientHandler {} impl ClientHandler for DummyClientHandler { - fn get_info(&self) -> ClientConfig { - ClientConfig::default() + fn get_info(&self) -> ClientInfo { + ClientInfo::default() } } diff --git a/crates/rmcp/tests/test_prompt_routers.rs b/crates/rmcp/tests/test_prompt_routers.rs index e730579..eecdb1e 100644 --- a/crates/rmcp/tests/test_prompt_routers.rs +++ b/crates/rmcp/tests/test_prompt_routers.rs @@ -20,6 +20,12 @@ 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 1453e0f..e91ecf9 100644 --- a/crates/rmcp/tests/test_protocol_version_negotiation.rs +++ b/crates/rmcp/tests/test_protocol_version_negotiation.rs @@ -1,40 +1,29 @@ //! Tests for protocol version negotiation in the default ServerHandler::initialize impl. //! -//! Handshake versions are echoed back; every other version falls back to one -//! the server can serve over `initialize`. +//! Known versions are echoed back; unknown versions fall back to LATEST. #![cfg(not(feature = "local"))] #![cfg(feature = "client")] -use std::{ - borrow::Cow, - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, -}; +use std::borrow::Cow; use rmcp::{ ClientHandler, ErrorData, RoleServer, ServerHandler, ServiceExt, - model::{ - ClientCapabilities, ClientConfig, ErrorCode, Implementation, InitializeRequestParams, - InitializeResult, ProtocolVersion, ServerConfig, - }, - service::{ClientInitializeError, RequestContext}, + model::{ClientInfo, InitializeRequestParams, InitializeResult, ProtocolVersion, ServerInfo}, + service::RequestContext, }; #[derive(Debug, Clone, Default)] struct EchoServer; impl ServerHandler for EchoServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::default() + fn get_info(&self) -> ServerInfo { + ServerInfo::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] = &[ +/// Every known version except `2026-07-28`, standing in for a server that has +/// not implemented that revision. +const NARROWED_VERSIONS: &[ProtocolVersion] = &[ ProtocolVersion::V_2024_11_05, ProtocolVersion::V_2025_03_26, ProtocolVersion::V_2025_06_18, @@ -45,30 +34,12 @@ const HANDSHAKE_VERSIONS: &[ProtocolVersion] = &[ struct NarrowedServer; impl ServerHandler for NarrowedServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::default() + fn get_info(&self) -> ServerInfo { + ServerInfo::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) + Cow::Borrowed(NARROWED_VERSIONS) } } @@ -79,12 +50,12 @@ impl ServerHandler for ModernOnlyServer { struct NarrowedOverridingServer; impl ServerHandler for NarrowedOverridingServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::default() + fn get_info(&self) -> ServerInfo { + ServerInfo::default() } fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { - Cow::Borrowed(HANDSHAKE_VERSIONS) + Cow::Borrowed(NARROWED_VERSIONS) } async fn initialize( @@ -102,8 +73,8 @@ struct VersionedClient { } impl ClientHandler for VersionedClient { - fn get_info(&self) -> ClientConfig { - let mut info = ClientConfig::default(); + fn get_info(&self) -> ClientInfo { + let mut info = ClientInfo::default(); info.protocol_version = self.protocol_version.clone(); info } @@ -117,28 +88,23 @@ 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 { - if let Ok(running) = server.serve(server_transport).await { - let _ = running.waiting().await; - } + let _ = server + .serve(server_transport) + .await + .expect("server should start") + .waiting() + .await; }); let client = VersionedClient { protocol_version: client_version, } .serve(client_transport) - .await?; + .await + .expect("client should connect"); let version = client .peer_info() @@ -147,55 +113,20 @@ async fn negotiate_with( .clone(); client.cancel().await.expect("client should cancel"); - Ok(version) + version } #[tokio::test] -async fn handshake_version_echoed_back() { - for version in HANDSHAKE_VERSIONS { +async fn known_version_echoed_back() { + for version in ProtocolVersion::KNOWN_VERSIONS { let negotiated = negotiated_version(version.clone()).await; assert_eq!( negotiated, *version, - "handshake version {version} should be echoed back" + "known 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(); @@ -209,7 +140,7 @@ async fn unknown_version_falls_back_to_latest() { #[tokio::test] async fn narrowed_server_still_echoes_versions_it_supports() { - for version in HANDSHAKE_VERSIONS { + for version in NARROWED_VERSIONS { let negotiated = negotiated_version_with(NarrowedServer, version.clone()).await; assert_eq!( negotiated, *version, @@ -238,91 +169,3 @@ async fn narrowed_server_caps_even_when_it_overrides_initialize() { "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 efae45f..44eb363 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, RoleClient, RoleServer, ServerHandler, ServiceError, + ClientHandler, RoleServer, ServerHandler, ServiceError, ServiceExt, model::{ - ClientConfig, ErrorCode, ErrorData, InitializeResult, ProtocolVersion, - ReadResourceRequestParams, ReadResourceResponse, + ClientInfo, ErrorCode, ErrorData, ProtocolVersion, ReadResourceRequestParams, + ReadResourceResponse, }, - service::{RequestContext, serve_directly}, + service::RequestContext, }; #[derive(Debug, Clone, Default)] @@ -33,40 +33,31 @@ struct VersionedClient { } impl ClientHandler for VersionedClient { - fn get_info(&self) -> ClientConfig { - let mut info = ClientConfig::default(); + fn get_info(&self) -> ClientInfo { + let mut info = ClientInfo::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 { - server.waiting().await?; + ResourceServer + .serve(server_transport) + .await? + .waiting() + .await?; anyhow::Ok(()) }); - let client = serve_directly::( - client_handler, - client_transport, - Some(server_peer_info.into()), - ); + let client = VersionedClient { + protocol_version: client_version, + } + .serve(client_transport) + .await + .expect("client should connect"); 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 61240ed..849a661 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, RoleClient, RoleServer, ServerHandler, + ClientHandler, RoleServer, ServerHandler, ServiceExt, model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, ClientConfig, ContentBlock, - ErrorData, InitializeResult, ProtocolVersion, ResultType, + CallToolRequestParams, CallToolResponse, CallToolResult, ClientInfo, ContentBlock, + ErrorData, ProtocolVersion, ResultType, }, - service::{RequestContext, serve_directly}, + service::RequestContext, }; #[derive(Debug, Clone, Default)] @@ -33,40 +33,27 @@ struct VersionedClient { } impl ClientHandler for VersionedClient { - fn get_info(&self) -> ClientConfig { - let mut info = ClientConfig::default(); + fn get_info(&self) -> ClientInfo { + let mut info = ClientInfo::default(); info.protocol_version = self.protocol_version.clone(); info } } -/// Wires the pair up directly on `client_version`. `2026-07-28` removed the -/// `initialize` handshake, so a peer on that revision is reached the way the -/// discover lifecycle leaves one: with the version already agreed. async fn call_tool_result_type(client_version: ProtocolVersion) -> Option { let (server_transport, client_transport) = tokio::io::duplex(4096); - let client_handler = VersionedClient { - protocol_version: client_version.clone(), - }; - let mut server_peer_info = InitializeResult::default(); - server_peer_info.protocol_version = client_version; - - let server = serve_directly::( - ToolServer, - server_transport, - Some(client_handler.get_info()), - ); let server_handle = tokio::spawn(async move { - server.waiting().await?; + ToolServer.serve(server_transport).await?.waiting().await?; anyhow::Ok(()) }); - let client = serve_directly::( - client_handler, - client_transport, - Some(server_peer_info.into()), - ); + let client = VersionedClient { + protocol_version: client_version, + } + .serve(client_transport) + .await + .expect("client should connect"); let result = client .call_tool(CallToolRequestParams::new("echo")) diff --git a/crates/rmcp/tests/test_sampling.rs b/crates/rmcp/tests/test_sampling.rs index 3c82fd4..b108ff4 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.5))] +#[case::float(serde_json::json!(3.14))] #[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 787d634..d4e20e1 100644 --- a/crates/rmcp/tests/test_sep_2260_request_association.rs +++ b/crates/rmcp/tests/test_sep_2260_request_association.rs @@ -1,19 +1,16 @@ #![cfg(all(feature = "server", feature = "client", not(feature = "local")))] -#![expect( - deprecated, - reason = "This test verifies request association for the deprecated sampling API" -)] +#![allow(deprecated)] use std::sync::{Arc, Mutex}; use rmcp::{ ClientHandler, RoleClient, RoleServer, ServerHandler, ServiceError, ServiceExt, model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, ClientConfig, ContentBlock, + CallToolRequestParams, CallToolResponse, CallToolResult, ClientInfo, ContentBlock, CreateMessageRequest, CreateMessageRequestParams, CreateMessageResult, ProtocolVersion, - SamplingMessage, ServerCapabilities, ServerConfig, ServerRequest, + SamplingMessage, ServerCapabilities, ServerInfo, ServerRequest, }, - service::{RequestContext, RunningService, serve_directly}, + service::RequestContext, }; use serde_json::{Value, json}; use tokio::{ @@ -21,16 +18,14 @@ use tokio::{ sync::oneshot, }; -type RequestResultSender = oneshot::Sender>; - #[derive(Clone)] struct SamplingServer { - outside: Arc>>, + outside: Arc>>>>, } impl ServerHandler for SamplingServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) } async fn call_tool( @@ -93,51 +88,28 @@ impl ClientHandler for SamplingClient { .with_stop_reason(CreateMessageResult::STOP_REASON_END_TURN)) } - fn get_info(&self) -> ClientConfig { - let mut info = ClientConfig::default(); + fn get_info(&self) -> ClientInfo { + let mut info = ClientInfo::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 { - running_server.waiting().await?; + let running = server.serve(server_transport).await?; + running.waiting().await?; anyhow::Ok(()) }); + let client = SamplingClient.serve(client_transport).await?; + let result = client .peer() .call_tool(CallToolRequestParams::new("sample")) @@ -157,16 +129,19 @@ 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 { - running_server.waiting().await?; + let running = server.serve(server_transport).await?; + running.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 index 72a0832..36595e0 100644 --- a/crates/rmcp/tests/test_sep_2260_stream_enforcement.rs +++ b/crates/rmcp/tests/test_sep_2260_stream_enforcement.rs @@ -26,8 +26,8 @@ use http::{HeaderName, HeaderValue}; use rmcp::{ ClientHandler, model::{ - ClientConfig, ClientJsonRpcMessage, CreateMessageRequestParams, CreateMessageResult, - ProtocolVersion, SamplingMessage, ServerCapabilities, ServerConfig, ServerJsonRpcMessage, + ClientInfo, ClientJsonRpcMessage, CreateMessageRequestParams, CreateMessageResult, + ProtocolVersion, SamplingMessage, ServerCapabilities, ServerInfo, ServerJsonRpcMessage, }, service::{ClientLifecycleMode, RequestContext, RoleClient, serve_client_with_lifecycle}, transport::streamable_http_client::{ @@ -80,7 +80,7 @@ impl StreamableHttpClient for ScriptedServer { // 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()); + let mut info = ServerInfo::new(ServerCapabilities::default()); info.protocol_version = ProtocolVersion::V_2026_07_28; let response = ServerJsonRpcMessage::response( rmcp::model::ServerResult::InitializeResult(info), @@ -152,8 +152,8 @@ impl ClientHandler for SamplingClient { )) } - fn get_info(&self) -> ClientConfig { - ClientConfig::default() + fn get_info(&self) -> ClientInfo { + ClientInfo::default() } } diff --git a/crates/rmcp/tests/test_sep_2260_stream_routing.rs b/crates/rmcp/tests/test_sep_2260_stream_routing.rs index e883ccd..824ef0c 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, ServerConfig, + ElicitationSchema, ServerCapabilities, ServerInfo, }, 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) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) } async fn call_tool( diff --git a/crates/rmcp/tests/test_server_discover_client.rs b/crates/rmcp/tests/test_server_discover_client.rs index 9046f6b..3b309dd 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, - ServerConfig, + ServerInfo, }, select_protocol_version, }; @@ -13,8 +13,8 @@ use rmcp::{ struct DiscoveryServer; impl ServerHandler for DiscoveryServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) .with_server_info(Implementation::new("discovery-server", "1.0.0")) } } diff --git a/crates/rmcp/tests/test_server_discover_http.rs b/crates/rmcp/tests/test_server_discover_http.rs index 7762c97..7a42f06 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, ServerConfig}, + model::{Implementation, ProtocolVersion, ServerCapabilities, ServerInfo}, 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) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) .with_server_info(Implementation::new("discovery-server", "1.0.0")) .with_instructions("Use the tools carefully") } diff --git a/crates/rmcp/tests/test_server_initialization.rs b/crates/rmcp/tests/test_server_initialization.rs index 5be178e..6542e29 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, ServerConfig, + ClientJsonRpcMessage, ProtocolVersion, ServerCapabilities, ServerInfo, 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) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 23ab8d0..e1e8852 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, ServerConfig, ToolsCapability}, + model::{Implementation, ServerCapabilities, ServerInfo, 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) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::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 cc2a4b5..02222ec 100644 --- a/crates/rmcp/tests/test_stateless_protocol_version.rs +++ b/crates/rmcp/tests/test_stateless_protocol_version.rs @@ -1,8 +1,7 @@ //! Tests for protocol version negotiation in stateless HTTP mode. //! -//! 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 +//! Supported versions are echoed back; unknown versions, and versions outside +//! the server's `supported_protocol_versions`, fall back to the handler's own //! version. #![cfg(not(feature = "local"))] @@ -11,8 +10,7 @@ use std::borrow::Cow; use rmcp::{ ErrorData, RoleServer, ServerHandler, model::{ - InitializeRequestParams, InitializeResult, ProtocolVersion, ServerCapabilities, - ServerConfig, + InitializeRequestParams, InitializeResult, ProtocolVersion, ServerCapabilities, ServerInfo, }, service::RequestContext, transport::streamable_http_server::{ @@ -25,8 +23,8 @@ use tokio_util::sync::CancellationToken; struct OverridingInitialize; impl ServerHandler for OverridingInitialize { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::default()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::default()) } async fn initialize( @@ -38,10 +36,9 @@ impl ServerHandler for OverridingInitialize { } } -/// 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] = &[ +/// Every known version except `2026-07-28`, standing in for a server that has +/// not implemented that revision. +const NARROWED_VERSIONS: &[ProtocolVersion] = &[ ProtocolVersion::V_2024_11_05, ProtocolVersion::V_2025_03_26, ProtocolVersion::V_2025_06_18, @@ -54,12 +51,12 @@ const HANDSHAKE_VERSIONS: &[ProtocolVersion] = &[ struct NarrowedOverridingInitialize; impl ServerHandler for NarrowedOverridingInitialize { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::default()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::default()) } fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { - Cow::Borrowed(HANDSHAKE_VERSIONS) + Cow::Borrowed(NARROWED_VERSIONS) } async fn initialize( @@ -151,15 +148,15 @@ async fn post_init(client: &reqwest::Client, url: &str, body_version: &str) -> s } #[tokio::test] -async fn stateless_json_init_echoes_handshake_versions_when_handler_overrides_initialize() { +async fn stateless_json_init_echoes_known_versions_when_handler_overrides_initialize() { let (client, url, ct) = spawn_server(stateless_json_config()).await; - for version in HANDSHAKE_VERSIONS { + for version in ProtocolVersion::KNOWN_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" + "known version {version} should be echoed back" ); } @@ -167,15 +164,15 @@ async fn stateless_json_init_echoes_handshake_versions_when_handler_overrides_in } #[tokio::test] -async fn stateless_sse_init_echoes_handshake_versions_when_handler_overrides_initialize() { +async fn stateless_sse_init_echoes_known_versions_when_handler_overrides_initialize() { let (client, url, ct) = spawn_server(stateless_sse_config()).await; - for version in HANDSHAKE_VERSIONS { + for version in ProtocolVersion::KNOWN_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" + "known version {version} should be echoed back" ); } @@ -201,7 +198,7 @@ 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 { + for version in NARROWED_VERSIONS { let resp = post_init(&client, &url, version.as_str()).await; assert_eq!( resp["result"]["protocolVersion"], diff --git a/crates/rmcp/tests/test_stateless_server_requests.rs b/crates/rmcp/tests/test_stateless_server_requests.rs index 9ab302b..4404e08 100644 --- a/crates/rmcp/tests/test_stateless_server_requests.rs +++ b/crates/rmcp/tests/test_stateless_server_requests.rs @@ -1,18 +1,14 @@ #![cfg(all(feature = "server", not(feature = "local")))] -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; +use std::sync::{Arc, Mutex}; use rmcp::{ ServerHandler, ServiceExt, model::{ ClientCapabilities, ClientJsonRpcMessage, ClientRequest, DiscoverRequest, DiscoverRequestParams, ErrorCode, ErrorData, Implementation, ListToolsRequest, - ListToolsResult, NumberOrString, PaginatedRequestParams, ProgressNotificationParam, - ProgressToken, ProtocolVersion, RequestId, RequestMetaObject, ServerJsonRpcMessage, - ServerNotification, + ListToolsResult, PaginatedRequestParams, ProtocolVersion, RequestId, RequestMetaObject, + ServerJsonRpcMessage, }, service::{MaybeSendFuture, RequestContext, RoleServer, ServerInitializeError}, transport::{IntoTransport, Transport}, @@ -166,7 +162,7 @@ async fn stateless_server_uses_each_requests_client_context() { } #[tokio::test] -async fn stateless_server_rejects_malformed_metadata_opener_with_error_response() { +async fn stateless_server_rejects_malformed_metadata_opener() { 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); @@ -190,24 +186,6 @@ async fn stateless_server_rejects_malformed_metadata_opener_with_error_response( )) .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"); }; @@ -216,73 +194,3 @@ async fn stateless_server_rejects_malformed_metadata_opener_with_error_response( 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 64be3fc..2e1e40a 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, - ServerConfig, + ServerInfo, }, }; use serde_json::{Value, json}; @@ -82,8 +82,8 @@ struct LargeResponseServer; impl ServerHandler for LargeResponseServer { #[allow(deprecated)] - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 13deeeb..ea49a41 100644 --- a/crates/rmcp/tests/test_streamable_http_4xx_error_body.rs +++ b/crates/rmcp/tests/test_streamable_http_4xx_error_body.rs @@ -7,10 +7,7 @@ use std::{collections::HashMap, sync::Arc}; use rmcp::{ - model::{ - ClientJsonRpcMessage, ClientRequest, DiscoverRequest, DiscoverRequestParams, PingRequest, - RequestId, - }, + model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId}, transport::streamable_http_client::{ StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse, }, @@ -48,13 +45,6 @@ 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. @@ -129,26 +119,3 @@ 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 deleted file mode 100644 index 2e5be84..0000000 --- a/crates/rmcp/tests/test_streamable_http_client_concurrency.rs +++ /dev/null @@ -1,1143 +0,0 @@ -//! 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 aa4ae63..553448e 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, ClientConfig, ServerCapabilities, ServerConfig}, + model::{CallToolRequestParams, ClientInfo, ServerCapabilities, ServerInfo}, 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) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 = ClientConfig::default().serve(transport).await?; + let client = ClientInfo::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 05efe21..31e7413 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, - ServerConfig, + ServerInfo, }, service::RequestContext, transport::streamable_http_server::{ @@ -37,8 +37,8 @@ struct CancelProbe { impl ServerHandler for CancelProbe { #[allow(deprecated)] - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 979ce7e..cc88f10 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, ServerConfig, + ProgressNotificationParam, ServerCapabilities, ServerInfo, }, service::RequestContext, transport::streamable_http_server::{ @@ -109,8 +109,8 @@ impl EventStore for InMemoryEventStore { struct ProgressServer; impl ServerHandler for ProgressServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) } async fn call_tool( diff --git a/crates/rmcp/tests/test_streamable_http_json_response.rs b/crates/rmcp/tests/test_streamable_http_json_response.rs index 62daa7a..b1c09f5 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, ServerConfig, + ProgressNotificationParam, ServerCapabilities, ServerInfo, }, 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) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 deleted file mode 100644 index c34e96e..0000000 --- a/crates/rmcp/tests/test_streamable_http_malformed_json_request.rs +++ /dev/null @@ -1,195 +0,0 @@ -#![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 76005a7..fcbeb41 100644 --- a/crates/rmcp/tests/test_streamable_http_protocol_version.rs +++ b/crates/rmcp/tests/test_streamable_http_protocol_version.rs @@ -87,17 +87,6 @@ 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) @@ -197,27 +186,6 @@ 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; @@ -250,37 +218,6 @@ async fn stateful_rejected_initial_posts_do_not_create_sessions() -> anyhow::Res 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; @@ -481,8 +418,8 @@ impl CountingServer { } impl ServerHandler for CountingServer { - fn get_info(&self) -> rmcp::model::ServerConfig { - rmcp::model::ServerConfig::new( + fn get_info(&self) -> rmcp::model::ServerInfo { + rmcp::model::ServerInfo::new( rmcp::model::ServerCapabilities::builder() .enable_tools() .build(), diff --git a/crates/rmcp/tests/test_streamable_http_stale_session.rs b/crates/rmcp/tests/test_streamable_http_stale_session.rs index c9ea214..be1ac02 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, ClientConfig, ClientJsonRpcMessage, ClientRequest, ErrorCode, + CallToolRequestParams, ClientInfo, 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 = ClientConfig::default().serve(transport).await?; + let mut client = ClientInfo::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 9e70805..510433e 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, ServerConfig, Tool}, + model::{ServerCapabilities, ServerInfo, 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) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 7359530..8ceef93 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, - ServerConfig, ServerNotification, ServerResult, SubscriptionFilter, + ServerInfo, ServerNotification, ServerResult, SubscriptionFilter, SubscriptionsAcknowledgedNotification, SubscriptionsAcknowledgedNotificationParams, SubscriptionsListenResult, }, @@ -44,8 +44,8 @@ impl ClientHandler for CountingClient { } impl ServerHandler for ToolsOnlyServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -80,8 +80,8 @@ impl ServerHandler for ToolsOnlyServer { struct ToolsAndPromptsServer; impl ServerHandler for ToolsAndPromptsServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -120,8 +120,8 @@ impl ServerHandler for ToolsAndPromptsServer { struct ResourceSubscriptionServer; impl ServerHandler for ResourceSubscriptionServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_resources() .enable_resources_subscribe() @@ -182,8 +182,8 @@ impl ServerHandler for RemoteCancellationServer { struct FloodServer; impl ServerHandler for FloodServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -217,8 +217,8 @@ struct ClosedSinkServer { struct LeakyServer; impl ServerHandler for LeakyServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -314,14 +314,14 @@ impl rmcp::service::Service for MalformedAcknowledgmentServer { Ok(()) } - fn get_info(&self) -> rmcp::model::ServerConfig { - ServerConfig::default() + fn get_info(&self) -> rmcp::model::ServerInfo { + ServerInfo::default() } } impl ServerHandler for ClosedSinkServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() diff --git a/crates/rmcp/tests/test_subscriptions_streamable_http.rs b/crates/rmcp/tests/test_subscriptions_streamable_http.rs index 70f102f..bfb9313 100644 --- a/crates/rmcp/tests/test_subscriptions_streamable_http.rs +++ b/crates/rmcp/tests/test_subscriptions_streamable_http.rs @@ -18,9 +18,8 @@ use std::{ use rmcp::{ ClientLifecycleMode, ClientServiceExt, ServerHandler, model::{ - ClientConfig, ClientRequest, Implementation, ListToolsRequest, ProtocolVersion, - RequestMetaObject, ServerCapabilities, ServerConfig, ServerNotification, - SubscriptionFilter, + ClientInfo, ClientRequest, Implementation, ListToolsRequest, ProtocolVersion, + RequestMetaObject, ServerCapabilities, ServerInfo, ServerNotification, SubscriptionFilter, }, service::{PeerRequestOptions, SubscriptionContext, SubscriptionEnd}, transport::{ @@ -53,8 +52,8 @@ impl ServerHandler for HttpSubscriptionServer { Cow::Borrowed(&[ProtocolVersion::V_2026_07_28, ProtocolVersion::V_2025_11_25]) } - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() @@ -160,7 +159,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 = ClientConfig::default() + let client = ClientInfo::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { @@ -205,7 +204,7 @@ async fn modern_http_graceful_close_returns_final_listen_result() -> anyhow::Res let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(url), ); - let client = ClientConfig::default() + let client = ClientInfo::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { @@ -241,7 +240,7 @@ async fn modern_http_stream_close_without_result_is_abrupt() -> anyhow::Result<( let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(url), ); - let client = ClientConfig::default() + let client = ClientInfo::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { @@ -274,7 +273,7 @@ async fn modern_http_lifecycle_stays_sessionless_for_older_application_version() let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(url), ); - let client = ClientConfig::default() + let client = ClientInfo::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { diff --git a/crates/rmcp/tests/test_task.rs b/crates/rmcp/tests/test_task.rs index 0178bab..ea1a259 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)] -struct SumArgs { - a: i32, - b: i32, +pub struct SumArgs { + pub a: i32, + pub b: i32, } #[derive(Clone)] @@ -107,8 +107,8 @@ impl ServerHandler for TaskServer { self.tasks.cancel_task(&request.task_id) } - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_tasks() @@ -117,8 +117,8 @@ impl ServerHandler for TaskServer { } } -fn tasks_client_info() -> ClientConfig { - ClientConfig::new( +fn tasks_client_info() -> ClientInfo { + ClientInfo::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) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::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 14a8757..cd87805 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, ServerConfig, Tool}, + model::{CallToolResponse, CallToolResult, ServerCapabilities, ServerInfo, Tool}, service::{MaybeSendFuture, NotificationContext}, }; use tokio::sync::{Notify, RwLock}; @@ -41,8 +41,8 @@ impl TestToolServer { } impl ServerHandler for TestToolServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 8642496..9b9530a 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, ClientConfig, ServerCapabilities, ServerConfig}, + model::{CallToolRequestParams, ClientInfo, ServerCapabilities, ServerInfo}, 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) -> ClientConfig { - ClientConfig::default() + fn get_info(&self) -> ClientInfo { + ClientInfo::default() } } @@ -549,8 +549,8 @@ impl ManualInfoServer { #[tool_handler] impl ServerHandler for ManualInfoServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_resources() @@ -573,47 +573,3 @@ 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 12a5d40..d2bbe86 100644 --- a/crates/rmcp/tests/test_tool_routers.rs +++ b/crates/rmcp/tests/test_tool_routers.rs @@ -21,6 +21,12 @@ 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 134bb1e..4c4ad52 100644 --- a/crates/rmcp/tests/test_unix_socket_transport.rs +++ b/crates/rmcp/tests/test_unix_socket_transport.rs @@ -13,12 +13,9 @@ use http::{HeaderName, HeaderValue}; use hyper_util::rt::TokioIo; use rmcp::{ ServiceExt, - model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId}, transport::{ StreamableHttpClientTransport, UnixSocketHttpClient, - streamable_http_client::{ - StreamableHttpClient, StreamableHttpClientTransportConfig, StreamableHttpError, - }, + streamable_http_client::StreamableHttpClientTransportConfig, }, }; use serde_json::json; @@ -38,10 +35,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") - && 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" { + if let Ok(v) = value.to_str() { + headers_map.insert(name_str.to_string(), v.to_string()); + } } } @@ -49,46 +46,46 @@ async fn mcp_handler( stored.extend(headers_map); drop(stored); - 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" + 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" + } } - } - }); - 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(), + ); + } } } @@ -299,56 +296,3 @@ 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 ffd3bf2..82a1fe9 100644 --- a/docs/OAUTH_SUPPORT.md +++ b/docs/OAUTH_SUPPORT.md @@ -15,7 +15,6 @@ 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 @@ -253,7 +252,7 @@ let transport = StreamableHttpClientTransport::with_client( ); // create client and connect to MCP server -let client_service = ClientConfig::default(); +let client_service = ClientInfo::default(); let client = client_service.serve(transport).await?; ``` @@ -295,130 +294,6 @@ 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) diff --git a/examples/clients/src/auth/client_credentials.rs b/examples/clients/src/auth/client_credentials.rs index e277b9e..55aa615 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::ClientConfig, + model::ClientInfo, transport::{ StreamableHttpClientTransport, auth::{AuthClient, ClientCredentialsConfig, OAuthState}, @@ -73,7 +73,7 @@ async fn main() -> Result<()> { ); // Connect to MCP server and list tools - let client_service = ClientConfig::default(); + let client_service = ClientInfo::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 8fa18df..1ab9313 100644 --- a/examples/clients/src/auth/oauth_client.rs +++ b/examples/clients/src/auth/oauth_client.rs @@ -9,7 +9,7 @@ use axum::{ }; use rmcp::{ RoleClient, ServiceExt, - model::ClientConfig, + model::ClientInfo, service::RunningService, transport::{ StreamableHttpClientTransport, @@ -58,7 +58,7 @@ async fn callback_handler( enum ConnectOutcome { /// The server accepted the unauthenticated connection. - Connected(RunningService), + Connected(RunningService), /// The server answered 401; authorize with this `WWW-Authenticate` /// challenge and reconnect. AuthRequired(String), @@ -72,7 +72,7 @@ async fn try_connect(http_client: reqwest::Client, server_url: &str) -> Result Ok(ConnectOutcome::Connected(client)), Err(error) => match error.auth_challenge() { Some(challenge) => Ok(ConnectOutcome::AuthRequired(challenge.to_string())), @@ -90,7 +90,7 @@ async fn authorize_and_connect( client_metadata_url: &str, code_receiver: oneshot::Receiver, output: &mut BufWriter, -) -> Result> { +) -> Result> { tracing::info!("Server requires authorization: {challenge}"); // initialize oauth state machine @@ -156,7 +156,7 @@ async fn authorize_and_connect( auth_client, StreamableHttpClientTransportConfig::with_uri(server_url), ); - Ok(ClientConfig::default().serve(transport).await?) + Ok(ClientInfo::default().serve(transport).await?) } #[tokio::main] diff --git a/examples/clients/src/progress_client.rs b/examples/clients/src/progress_client.rs index 3990006..a9f68c1 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, ClientConfig, Implementation, + CallToolRequestParams, ClientCapabilities, ClientInfo, Implementation, ProgressNotificationParam, }, service::{NotificationContext, RoleClient}, @@ -100,10 +100,10 @@ impl ProgressAwareClient { } fn stop_tracking(&self) { - if let Ok(mut tracker_opt) = self.tracker.lock() - && let Some(tracker) = tracker_opt.take() - { - tracker.print_summary(); + if let Ok(mut tracker_opt) = self.tracker.lock() { + if 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() - && let Some(tracker) = tracker_opt.as_ref() - { - tracker.handle_progress(¶ms); + if let Ok(tracker_opt) = self.tracker.lock() { + if let Some(tracker) = tracker_opt.as_ref() { + tracker.handle_progress(¶ms); + } } } - fn get_info(&self) -> ClientConfig { - ClientConfig::new( + fn get_info(&self) -> ClientInfo { + ClientInfo::new( ClientCapabilities::default(), Implementation::new("progress-test-client", "1.0.0"), ) @@ -182,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() - && let Some(text) = content.as_text() - { - tracing::info!("Processing completed: {}", text.text); + if let Some(content) = tool_result.content.first() { + if let Some(text) = content.as_text() { + tracing::info!("Processing completed: {}", text.text); + } } service.cancel().await?; @@ -236,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() - && let Some(text) = content.as_text() - { - tracing::info!("processing completed: {}", text.text); + if let Some(content) = tool_result.content.first() { + if 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 1079586..cc7c5f1 100644 --- a/examples/clients/src/sampling_stdio.rs +++ b/examples/clients/src/sampling_stdio.rs @@ -1,8 +1,3 @@ -#![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 f841a87..0c27fd3 100644 --- a/examples/clients/src/streamable_http.rs +++ b/examples/clients/src/streamable_http.rs @@ -1,15 +1,14 @@ use anyhow::Result; use rmcp::{ - ClientLifecycleMode, ClientServiceExt, - model::{ - CallToolRequestParams, ClientCapabilities, ClientConfig, Implementation, ProtocolVersion, - }, + ServiceExt, + model::{CallToolRequestParams, ClientCapabilities, ClientInfo, Implementation}, 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() @@ -18,31 +17,27 @@ async fn main() -> Result<()> { .with(tracing_subscriber::fmt::layer()) .init(); let transport = StreamableHttpClientTransport::from_uri("http://localhost:8000/mcp"); - let client_info = ClientConfig::new( + let client_info = ClientInfo::new( ClientCapabilities::default(), - Implementation::new("streamable-http-client", "0.0.1"), + Implementation::new("test sse client", "0.0.1"), ); - 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); - })?; + let client = client_info.serve(transport).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")) + .call_tool( + CallToolRequestParams::new("increment") + .with_arguments(serde_json::json!({}).as_object().cloned().unwrap()), + ) .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 9c04531..7ab612b 100644 --- a/examples/clients/src/subscriptions_streamhttp.rs +++ b/examples/clients/src/subscriptions_streamhttp.rs @@ -1,6 +1,6 @@ use rmcp::{ ClientLifecycleMode, ClientServiceExt, - model::{ClientConfig, ProtocolVersion, SubscriptionFilter}, + model::{ClientInfo, 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 = ClientConfig::default() + let client = ClientInfo::default() .serve_with_lifecycle( transport, ClientLifecycleMode::Discover { diff --git a/examples/clients/src/task_stdio.rs b/examples/clients/src/task_stdio.rs index a946f75..465cb70 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::ClientConfig::new( + let client_info = rmcp::model::ClientInfo::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 a812977..6a634d8 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 ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::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,8 +278,7 @@ impl ServerHandler for Counter { let initialize_uri = &http_request_part.uri; tracing::info!(?initialize_headers, %initialize_uri, "initialize from http server"); } - context.peer.set_peer_info(request.clone()); - self.negotiate_initialize(&request) + Ok(self.get_info()) } } diff --git a/examples/servers/src/common/progress_demo.rs b/examples/servers/src/common/progress_demo.rs index 68a32ec..253dbc2 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("No progress token", None)); + return Err(McpError::internal_error(format!("No progress token"), None)); }; let Ok(progress_token) = serde_json::from_value::(progress_token.clone()) else { return Err(McpError::internal_error( - "Invalid format of the progress token", + format!("Invalid format of the progress token"), None, )); }; @@ -131,8 +131,8 @@ impl ProgressDemo { #[tool_handler] impl ServerHandler for ProgressDemo { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 30aed17..4d8d53a 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) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_tasks() diff --git a/examples/servers/src/completion_stdio.rs b/examples/servers/src/completion_stdio.rs index c7894ea..812ed31 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() - && let Some(first) = candidate.chars().next() - { - first_chars.push(first.to_lowercase().next().unwrap_or('\0')); + 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')); + } } } // 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 - && 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() == 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() != first_chars.len() { @@ -193,7 +193,7 @@ impl SqlQueryServer { } } -#[prompt_router(router = "prompt_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(router = self.prompt_router)] +#[prompt_handler] impl ServerHandler for SqlQueryServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::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 bf763d3..34c4b15 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) - && client.redirect_uri == redirect_uri - { - return Some(client.clone()); + if let Some(client) = clients.get(client_id) { + if client.redirect_uri.contains(&redirect_uri.to_string()) { + return Some(client.clone()); + } } None } diff --git a/examples/servers/src/elicitation_enum_inference.rs b/examples/servers/src/elicitation_enum_inference.rs index e565892..bed5c1d 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(router = tool_router)] +#[tool_router] impl ElicitationEnumFormServer { pub fn new() -> Self { Self { @@ -153,10 +153,10 @@ impl ElicitationEnumFormServer { } } -#[tool_handler(router = self.tool_router)] +#[tool_handler] impl ServerHandler for ElicitationEnumFormServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 5d7f280..d506a9c 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(router = tool_router)] +#[tool_router] impl ElicitationServer { #[tool(description = "Greet user with name collection")] async fn greet_user( @@ -145,10 +145,10 @@ impl ElicitationServer { } } -#[tool_handler(router = self.tool_router)] +#[tool_handler] impl ServerHandler for ElicitationServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 b6c87b6..40b53c3 100644 --- a/examples/servers/src/mrtr.rs +++ b/examples/servers/src/mrtr.rs @@ -38,83 +38,30 @@ use rmcp::{ model::*, service::{RequestContext, RoleClient, RoleServer}, }; -use serde::{Deserialize, Serialize}; use serde_json::json; -/// A fixed key for this example. Real deployments should load at least 32 bytes -/// of random secret key material from a secret manager and keep it stable for +/// 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 /// the lifetime of any in-flight MRTR exchange. const REQUEST_STATE_KEY: &[u8] = b"example-request-state-signing-key-32b!"; -#[derive(Debug, Serialize, Deserialize)] -struct WeatherRequestState { - awaiting: WeatherInput, -} - -#[derive(Debug, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -enum WeatherInput { - City, -} - /// A server that needs a city name before it can answer a weather query. #[derive(Clone)] struct WeatherServer { codec: RequestStateCodec, } -impl WeatherServer { - fn new(key: impl Into>) -> Result { - Ok(Self { - codec: RequestStateCodec::try_new(key)?, - }) - } -} - -impl WeatherServer { - fn open_request_state(&self, sealed: &str) -> Result { - self.codec - .open_json(sealed) - .map_err(|_| ErrorData::invalid_params("tampered or unknown request state", None)) - } -} - -fn finish_weather_request( - input_responses: Option<&InputResponses>, -) -> Result { - let response = input_responses - .and_then(|responses| responses.get("city")) - .ok_or_else(|| ErrorData::invalid_params("missing city elicitation response", None))?; - let response = ElicitResult::deserialize(response) - .map_err(|_| ErrorData::invalid_params("invalid city elicitation response", None))?; - - let message = match response.action { - ElicitationAction::Accept => { - let city = response - .content - .as_ref() - .and_then(|content| content["city"].as_str()) - .ok_or_else(|| { - ErrorData::invalid_params("accepted elicitation did not include a city", None) - })?; - format!("It is sunny in {city}.") +impl Default for WeatherServer { + fn default() -> Self { + Self { + codec: RequestStateCodec::new(REQUEST_STATE_KEY), } - ElicitationAction::Decline => "I cannot look up the weather without a city.".into(), - ElicitationAction::Cancel => "Weather lookup cancelled.".into(), - _ => { - return Err(ErrorData::invalid_params( - "unsupported elicitation action", - None, - )); - } - }; - - Ok(CallToolResult::success(vec![ContentBlock::text(message)])) + } } impl ServerHandler for WeatherServer { - fn get_info(&self) -> ServerConfig { - let mut info = ServerConfig::new(ServerCapabilities::builder().enable_tools().build()); + fn get_info(&self) -> ServerInfo { + let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build()); // MRTR requires 2026-07-28 or newer. info.protocol_version = ProtocolVersion::V_2026_07_28; info @@ -131,9 +78,7 @@ impl ServerHandler for WeatherServer { None => { let sealed = self .codec - .seal_json(&WeatherRequestState { - awaiting: WeatherInput::City, - }) + .seal_json(&json!({ "awaiting": "city" })) .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; let mut input_requests = InputRequests::new(); @@ -158,11 +103,21 @@ 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 = self.open_request_state(&sealed)?; - match state.awaiting { - WeatherInput::City => finish_weather_request(request.input_responses.as_ref()) - .map(CallToolResponse::from), - } + 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()) } } } @@ -173,8 +128,8 @@ impl ServerHandler for WeatherServer { struct InteractiveClient; impl ClientHandler for InteractiveClient { - fn get_info(&self) -> ClientConfig { - ClientConfig::new( + fn get_info(&self) -> ClientInfo { + ClientInfo::new( ClientCapabilities::builder().enable_elicitation().build(), Implementation::new("mrtr-example-client", env!("CARGO_PKG_VERSION")), ) @@ -197,11 +152,10 @@ 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 = server + let server = WeatherServer::default() .serve(server_transport) .await .expect("server should start"); @@ -209,7 +163,7 @@ async fn main() -> anyhow::Result<()> { }); // Connect the client (this performs the initialize handshake). - let client = InteractiveClient.serve(client_transport).await?; + let client = InteractiveClient::default().serve(client_transport).await?; // 1. High-level auto mode: the SDK fulfils the elicitation and retries for us. println!("== auto mode (call_tool) =="); @@ -243,91 +197,3 @@ 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 7f018d0..7b6e285 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(router = "prompt_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 - && !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 { + 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.", + )); + } } messages.push(PromptMessage::new_text( @@ -361,10 +361,10 @@ impl PromptServer { } } -#[prompt_handler(router = self.prompt_router)] +#[prompt_handler] impl ServerHandler for PromptServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_prompts().build()).with_instructions( + fn get_info(&self) -> ServerInfo { + ServerInfo::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 41f74c7..be230ad 100644 --- a/examples/servers/src/sampling_stdio.rs +++ b/examples/servers/src/sampling_stdio.rs @@ -1,7 +1,4 @@ -#![expect( - deprecated, - reason = "This example demonstrates the deprecated MCP sampling API" -)] +#![allow(deprecated)] use std::sync::Arc; use anyhow::Result; @@ -21,8 +18,8 @@ use tracing_subscriber::{self, EnvFilter}; pub struct SamplingDemoServer; impl ServerHandler for SamplingDemoServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + fn get_info(&self) -> ServerInfo { + ServerInfo::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 7615d59..aead85f 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, ServerConfig, SubscriptionFilter}, + model::{ProtocolVersion, ServerCapabilities, ServerInfo, 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) -> ServerConfig { - ServerConfig::new( + fn get_info(&self) -> ServerInfo { + ServerInfo::new( ServerCapabilities::builder() .enable_tools() .enable_tool_list_changed() diff --git a/scripts/keycloak-oauth-fixture.sh b/scripts/keycloak-oauth-fixture.sh deleted file mode 100755 index 50c9b85..0000000 --- a/scripts/keycloak-oauth-fixture.sh +++ /dev/null @@ -1,56 +0,0 @@ -#!/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"