diff --git a/.agents/skills/fix-bug/SKILLS.md b/.agents/skills/fix-bug/SKILLS.md new file mode 100644 index 0000000..6956a54 --- /dev/null +++ b/.agents/skills/fix-bug/SKILLS.md @@ -0,0 +1,13 @@ +--- +name: fix-bug +description: Confirm, debug and fix bugs reported via GitHub issues. +--- + +# Step-by-step instructions + +1. Use the `gh` CLI to fetch information about the issue. +2. Create a reproduction test case. +3. Create an implementation plan that describes how you intend to fix the bug. Don't proceed without approval. +4. Implement the fix. +5. Verify the fix with the test case and run all tests to ensure no regressions. +6. Provide a summary of the changes. diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..ce398c1 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,25 @@ +--- +name: Bug report +about: Create a report to help us improve +title: '' +labels: '' +assignees: '' +type: Bug + +--- + +**Describe the bug** +A clear and concise description of what the bug is. + +**To Reproduce** +Steps to reproduce the behavior: +1. + +**Expected behavior** +A clear and concise description of what you expected to happen. + +**Logs** +If applicable, add logs to help explain your problem. + +**Additional context** +Add any other context about the problem here. diff --git a/.github/ISSUE_TEMPLATE/enhancement.md b/.github/ISSUE_TEMPLATE/enhancement.md new file mode 100644 index 0000000..d2ecf10 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/enhancement.md @@ -0,0 +1,21 @@ +--- +name: Enhancement +about: Suggest an idea for this project +title: '' +labels: '' +assignees: '' +type: Enhancement + +--- + +**Is your feature request related to a problem? Please describe.** +A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] + +**Describe the solution you'd like** +A clear and concise description of what you want to happen. + +**Describe alternatives you've considered** +A clear and concise description of any alternative solutions or features you've considered. + +**Additional context** +Add any other context or screenshots about the feature request here. diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..f9fbfe1 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,15 @@ +version: 2 +updates: +# Check for updates to GitHub Actions every month. +- package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "monthly" + # Group all github/codeql-action/* sub-paths (init, analyze, ...) into a + # single PR so they are always bumped together. Bumping them independently + # leaves the workflow with mismatched versions across steps, which + # github/codeql-action rejects at runtime. + groups: + codeql-action: + patterns: + - "github/codeql-action/*" diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml new file mode 100644 index 0000000..c4aa91c --- /dev/null +++ b/.github/workflows/codeql.yml @@ -0,0 +1,48 @@ +name: "CodeQL Advanced" + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + schedule: + - cron: '31 9 * * 4' + +# Declare default permissions as read only. +permissions: read-all + +jobs: + analyze: + name: Analyze (${{ matrix.language }}) + runs-on: ubuntu-latest + permissions: + # required for all workflows + security-events: write + + # required to fetch internal or private CodeQL packs + packages: read + + # only required for workflows in private repositories + actions: read + contents: read + + strategy: + fail-fast: false + matrix: + include: + - language: actions + build-mode: none + - language: go + build-mode: autobuild + steps: + - name: Checkout repository + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - name: Initialize CodeQL + uses: github/codeql-action/init@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9 + with: + languages: ${{ matrix.language }} + build-mode: ${{ matrix.build-mode }} + - name: Perform CodeQL Analysis + uses: github/codeql-action/analyze@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9 + with: + category: "/language:${{matrix.language}}" diff --git a/.github/workflows/conformance.yml b/.github/workflows/conformance.yml new file mode 100644 index 0000000..6f54df4 --- /dev/null +++ b/.github/workflows/conformance.yml @@ -0,0 +1,74 @@ +name: Conformance Tests + +on: + push: + branches: [main] + pull_request: + workflow_dispatch: + +concurrency: + group: conformance-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +env: + CONFORMANCE_VERSION: "0.2.0-alpha.11" + +jobs: + server-conformance: + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - name: Set up Go + uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 + with: + go-version: "^1.26" + - name: Set up Node.js + uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 + with: + node-version: 22 + - name: Build everything-server + run: go build -o everything-server ./conformance/everything-server + - name: Run conformance tests (2025-11-25, stateful) + run: | + ./everything-server -http=":3001" -stateless=false & + SERVER_PID=$! + trap "kill -9 $SERVER_PID 2>/dev/null || true; wait $SERVER_PID 2>/dev/null || true" EXIT + timeout 30 bash -c 'until curl -s http://localhost:3001/mcp > /dev/null; do sleep 0.5; done' + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" server \ + --url http://localhost:3001/mcp \ + --requirements 2025-11-25 \ + --expected-failures ./conformance/baseline.yml + - name: Run conformance tests (2026-07-28, stateless) + run: | + ./everything-server -http=":3001" -stateless & + SERVER_PID=$! + trap "kill -9 $SERVER_PID 2>/dev/null || true; wait $SERVER_PID 2>/dev/null || true" EXIT + timeout 30 bash -c 'until curl -s http://localhost:3001/mcp > /dev/null; do sleep 0.5; done' + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" server \ + --url http://localhost:3001/mcp \ + --requirements 2026-07-28 \ + --expected-failures ./conformance/baseline.yml + + client-conformance: + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - name: Set up Go + uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 + with: + go-version: "^1.26" + - name: Set up Node.js + uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 + with: + node-version: 22 + - name: Run conformance tests (full suite) + run: | + npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" client \ + --command "go run ./conformance/everything-client" \ + --suite all \ + --expected-failures ./conformance/baseline.yml diff --git a/.github/workflows/docs-check.yml b/.github/workflows/docs-check.yml index 984e878..d255e63 100644 --- a/.github/workflows/docs-check.yml +++ b/.github/workflows/docs-check.yml @@ -16,9 +16,11 @@ jobs: runs-on: ubuntu-latest steps: - name: Set up Go - uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5 + uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 + with: + go-version: "^1.26" - name: Check out code - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Check docs are up-to-date run: | go generate ./... diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml new file mode 100644 index 0000000..2daa939 --- /dev/null +++ b/.github/workflows/nightly.yml @@ -0,0 +1,40 @@ +name: Nightly + +on: + schedule: + - cron: '0 2 * * *' # 2 AM UTC daily + + workflow_dispatch: + +permissions: + contents: read + +jobs: + vulncheck: + # Don't run scheduled builds on forks. + if: github.event_name != 'schedule' || github.repository == 'modelcontextprotocol/go-sdk' + runs-on: ubuntu-latest + steps: + - name: Run govulncheck + uses: golang/govulncheck-action@032d45514ae346b1db93c04b0c90b841c370344f # v1.1.0 + with: + go-version-file: go.mod + go-package: ./... + + conformance: + # Don't run scheduled builds on forks. + if: github.event_name != 'schedule' || github.repository == 'modelcontextprotocol/go-sdk' + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - name: Set up Go + uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 + with: + go-version: "^1.26" + - name: Set up Node.js + uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 + with: + node-version: "22" + - name: Run server conformance tests + run: ./scripts/server-conformance.sh diff --git a/.github/workflows/publish-docs.yml b/.github/workflows/publish-docs.yml new file mode 100644 index 0000000..4d6d31e --- /dev/null +++ b/.github/workflows/publish-docs.yml @@ -0,0 +1,32 @@ +name: Publish Docs + +on: + release: + types: [published] + workflow_dispatch: + +permissions: read-all + +jobs: + deploy: + permissions: + contents: write + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - name: Configure Git Credentials + run: | + git config user.name github-actions[bot] + git config user.email 41898282+github-actions[bot]@users.noreply.github.com + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: 3.x + - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV + - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + key: mkdocs-material-${{ env.cache_id }} + path: .cache + restore-keys: | + mkdocs-material- + - run: pip install mkdocs-material + - run: mkdocs gh-deploy --force diff --git a/.github/workflows/scorecard.yml b/.github/workflows/scorecard.yml new file mode 100644 index 0000000..98db4c4 --- /dev/null +++ b/.github/workflows/scorecard.yml @@ -0,0 +1,78 @@ +# This workflow uses actions that are not certified by GitHub. They are provided +# by a third-party and are governed by separate terms of service, privacy +# policy, and support documentation. + +name: Scorecard supply-chain security +on: + # For Branch-Protection check. Only the default branch is supported. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection + branch_protection_rule: + # To guarantee Maintained check is occasionally updated. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained + schedule: + - cron: '41 23 * * 5' + push: + branches: [ "main" ] + +# Declare default permissions as read only. +permissions: read-all + +jobs: + analysis: + name: Scorecard analysis + runs-on: ubuntu-latest + # `publish_results: true` only works when run from the default branch. conditional can be removed if disabled. + if: github.event.repository.default_branch == github.ref_name || github.event_name == 'pull_request' + permissions: + # Needed to upload the results to code-scanning dashboard. + security-events: write + # Needed to publish results and get a badge (see publish_results below). + id-token: write + # Uncomment the permissions below if installing in a private repository. + # contents: read + # actions: read + + steps: + - name: "Checkout code" + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + + - name: "Run analysis" + uses: ossf/scorecard-action@2d1146689b8cda280b9bc96326124645441f03bc # v2.4.4 + with: + results_file: results.sarif + results_format: sarif + # (Optional) "write" PAT token. Uncomment the `repo_token` line below if: + # - you want to enable the Branch-Protection check on a *public* repository, or + # - you are installing Scorecard on a *private* repository + # To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional. + # repo_token: ${{ secrets.SCORECARD_TOKEN }} + + # Public repositories: + # - Publish results to OpenSSF REST API for easy access by consumers + # - Allows the repository to include the Scorecard badge. + # - See https://github.com/ossf/scorecard-action#publishing-results. + # For private repositories: + # - `publish_results` will always be set to `false`, regardless + # of the value entered here. + publish_results: true + + # (Optional) Uncomment file_mode if you have a .gitattributes with files marked export-ignore + # file_mode: git + + # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF + # format to the repository Actions tab. + - name: "Upload artifact" + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: SARIF file + path: results.sarif + retention-days: 5 + + # Upload the results to GitHub's code scanning dashboard (optional). + # Commenting out will disable upload of results to your repo's Code Scanning dashboard + - name: "Upload to code-scanning" + uses: github/codeql-action/upload-sarif@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9 + with: + sarif_file: results.sarif diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 68a3c42..5c351b3 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -14,11 +14,11 @@ jobs: runs-on: ubuntu-latest steps: - name: Check out code - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Set up Go - uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5 + uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 with: - go-version: "^1.25" + go-version: "^1.26" - name: Check formatting run: | unformatted=$(gofmt -l .) @@ -31,7 +31,7 @@ jobs: - name: Run Go vet run: go vet ./... - name: Run staticcheck - uses: dominikh/staticcheck-action@024238d2898c874f26d723e7d0ff4308c35589a2 # v1 + uses: dominikh/staticcheck-action@9716614d4101e79b4340dd97b10e54d68234e431 # v1.4.1 with: install-go: false version: "v0.6.1" @@ -40,12 +40,12 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - go: ["1.23", "1.24", "1.25"] + go: ["1.25", "1.26"] steps: - name: Check out code - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Set up Go - uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5 + uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 with: go-version: ${{ matrix.go }} - name: Test @@ -55,10 +55,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Check out code - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Set up Go - uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5 + uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 with: - go-version: "1.24" + go-version: "1.26" - name: Test with -race run: go test -v -race ./... diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..bc65991 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,58 @@ +# AGENTS.md + +## Project Overview + +This repository contains the official Go SDK for the Model Context Protocol (MCP). +The SDK is designed to be idiomatic, future-proof, and extensible. + +### Key Packages + +- `mcp`: The core package defining the primary APIs for constructing and using MCP clients and servers. This is where most logic resides. +- `jsonrpc`: Provides the JSON-RPC 2.0 transport layer. Use this if implementing custom transports. +- `auth`: Primitives for supporting OAuth. +- `oauthex`: Extensions to the OAuth protocol, such as Protected Resource Metadata. +- `internal`: Internal implementation details not exposed to users. +- `examples`: Example clients and servers. Use these as references for usage patterns. + +## Development Setup + +The project uses the standard Go toolchain. + +- **Build**: `go build ./...` +- **Test**: `go test ./...` + +## Testing + +- **Unit Tests**: Run `go test ./...` to run all unit tests. +- **Conformance Tests**: Use the following scripts to run the official MCP conformance tests against the SDK. + - `./scripts/server-conformance.sh` for server tests. + - `./scripts/client-conformance.sh` for client tests. + - The scripts download the latest conformance suite from npm by default. + - To get possible options pass the `--help` flag to the script. + +## Development Guidelines + +### Code Style + +- Follow standard Go conventions (Effective Go). +- Use `gofmt` to format code. +- Add copyright headers to all new Go files: + ```go + // Copyright 2025 The Go MCP SDK Authors. All rights reserved. + // Use of this source code is governed by the license + // that can be found in the LICENSE file. + ``` +- Do not add comments to the code unless they are really necessary: + - Prefer self-documenting code. + - Focus on the "why" not the "what" in comments. + +### Documentation + +- **README.md**: Do NOT edit `README.md` directly. It is generated from `internal/readme/README.src.md`. + - Edit `internal/readme/README.src.md`. + - Run `go generate ./internal/readme` to regenerate. + - Commit both files. +- **docs/**: Do NOT edit `docs/` directory directly. It is generated from files in `internal/docs`. + - Edit `internal/docs/*.src.md`. + - Run `go generate ./internal/docs` to regenerate. + - Commit files from both directories. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 001c35d..3d8b090 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,3 +1,4 @@ + # Contributing to the Go MCP SDK Thank you for your interest in contributing! The Go SDK needs active @@ -6,11 +7,31 @@ new and emerging use-cases. We welcome all forms of contribution, from filing and reviewing issues, to contributing fixes, to proposing and implementing new features. -As described in the [design document](./design/design.md), it is important for +As described in the [design document](https://github.com/modelcontextprotocol/go-sdk/blob/main/design/design.md), it is important for the MCP SDK to remain idiomatic, future-proof, and extensible. The process described here is intended to ensure that the SDK evolves safely and transparently, while adhering to these goals. +1. [Development setup](#development-setup) + 1. [Conformance tests](#conformance-tests) +1. [Filing issues](#filing-issues) + 1. [Bugs](#bugs) + 1. [Proposals](#proposals) + 1. [Design discussion](#design-discussion) +1. [Contributing code](#contributing-code) + 1. [Adding and updating dependencies](#adding-and-updating-dependencies) + 1. [Updating the README](#updating-the-readme) +1. [Versioning](#versioning) + 1. [What counts as a breaking change](#what-counts-as-a-breaking-change) + 1. [Deprecation](#deprecation) +1. [Timeouts](#timeouts) +1. [Code of conduct](#code-of-conduct) +1. [Governance](#governance) + 1. [Working Group meetings](#working-group-meetings) + 1. [Discord](#discord) + 1. [Antitrust considerations](#antitrust-considerations) + + ## Development setup This module can be built and tested using the standard Go toolchain. Run `go @@ -26,6 +47,34 @@ containing the SDK, run: go work init ./project ./go-sdk ``` +### Conformance tests + +The SDK includes a script to run the official MCP conformance tests against the +SDK's conformance server: + +```sh +./scripts/conformance.sh +``` + +By default, results are cleaned up after the script runs. To save results to a +specific directory: + +```sh +./scripts/conformance.sh --result_dir ./conformance-results +``` + +To run against a local checkout of the +[conformance repo](https://github.com/modelcontextprotocol/conformance) instead +of the latest npm release: + +```sh +./scripts/conformance.sh --conformance_repo ~/src/conformance +``` + +Note: you must run `npm install` in the conformance repo first. + +Run `./scripts/conformance.sh --help` for more options. + ## Filing issues This project uses the [GitHub issue @@ -95,14 +144,38 @@ Changes should be high quality and well tested, and should generally follow the messages should follow the [format used by the Go project](https://go.dev/wiki/CommitMessage). -Unless otherwise noted, the Go source files are distributed under the MIT-style -license found in the LICENSE file. All Go files in the SDK should have a -copyright header following the format below: +Unless otherwise noted, the Go source files are distributed under the license +found in the LICENSE file. New contributions are licensed under Apache 2.0. All +Go files in the SDK should have a copyright header following the format below: ```go // Copyright 2025 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. +``` + +### Adding and updating dependencies + +In general, the SDK tries to use as few dependencies as possible. Each new +dependency is a potential source for bugs, churn, and conflicts for our users. +Therefore, we require a [proposal](#proposals) for any new module dependency, +including upgrading an existing module to a new major version. New dependencies +should be evaluated for their stability and security, and should be +well-established in the Go ecosystem. + +In general, dependencies should be for internal use by the SDK implementation, +or for testing. Do not include types from dependencies in the SDK API. + +On the other hand, updating existing dependencies can be done at any time +without a proposal, as long as their major version does not change. Prefer to +update dependencies immediately following a release of the SDK, to allow as +much time as possible to find issues with the new version. + +After any change to dependencies, run govulncheck to check them for +vulnerabilities. + +``` +go run golang.org/x/vuln/cmd/govulncheck@latest ``` ### Updating the README @@ -118,6 +191,57 @@ The CI system will automatically check that the README is up-to-date by running `go generate ./internal/readme` and verifying no changes result. If you see a CI failure about the README being out of sync, follow the steps above to regenerate it. +## Versioning + +The SDK follows [semantic versioning](https://semver.org), as the Go module +system requires: the major version is part of the module path, so a breaking +release would import as `github.com/modelcontextprotocol/go-sdk/v2`. A breaking +change therefore cannot reach existing users by accident; they have to change +their import path to receive one. + +This policy covers the exported API of the SDK's importable packages — `mcp`, +`jsonrpc`, `auth`, `auth/extauth` and `oauthex`. Everything under `internal/` +is not importable outside the module and may change in any release. + +Which MCP spec revisions each SDK version speaks is documented in the +[README](README.md#version-compatibility). Support for a new revision arrives +in a minor release; dropping a revision the SDK previously negotiated is a +breaking change. + +### What counts as a breaking change + +- Removing or renaming an exported identifier. +- Changing the signature of an exported function or method. +- Removing or renaming a field of an exported struct, or changing its type. +- Adding a method to an exported interface that users are meant to implement. +- Changing documented behavior that callers could reasonably rely on. +- Dropping support for an MCP spec revision. + +The following are not breaking, and may land in a minor or patch release: + +- Adding an exported identifier, or a field to an exported struct. +- Fixing behavior that contradicted its documentation. +- Refactoring that leaves the exported API unchanged. +- Adding support for a new MCP spec revision. +- Raising the minimum Go version to one still supported upstream, per the + [README](README.md#version-compatibility). + +### Deprecation + +An API on its way out is marked with a `// Deprecated:` comment naming its +replacement, which `gopls` and `staticcheck` surface at call sites. It keeps +working until the next major version, since removing it would itself be a +breaking change. + +Protocol features that the *specification* deprecates are a separate matter and +follow the spec's timeline rather than this policy. The features deprecated by +[SEP-2577](https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging) +— roots, sampling and logging — remain supported for at least twelve months, as +noted in the [README](README.md#version-compatibility). + +Any change to the exported API, breaking or not, goes through the +[proposal process](#proposals). + ## Timeouts If a contributor hasn't responded to issue questions or PR comments in two weeks, @@ -159,134 +283,9 @@ GitHub discussions, or public steering meetings. ### Antitrust considerations -It is important that the SDK avoids bias toward specific integration paths or -providers. The antitrust policy below details precise terms that should be -adhered to in order to ensure that the evolution and governance of the SDK remain -procompetitive. - -### Antitrust policy - -Note: all changes to the policy in this section must be approved by legal -counsel. - -The Go+Anthropic MCP SDK Working Group (the “Working Group”) is established to -guide and review technical contributions to an open-source Go software -development kit (“SDK”) for the Model Context Protocol (“MCP”). The Working Group -and its members are committed to operating for procompetitive purposes that -accelerate AI development and benefit businesses and consumers. This -collaboration is focused on technical and infrastructure objectives – namely, -developing and maintaining a neutral, open-source, and MIT-licensed library. -Google and Anthropic, as well as other stakeholders, participate with the -understanding that the Working Group’s sole purpose is to improve interoperability -and innovation in the MCP ecosystem. - -Antitrust law recognizes that when competitors collaborate for valid reasons -(e.g., joint R&D or standard-setting), such collaborations can be -procompetitive. This Antitrust Compliance Policy (the “Policy”) therefore -outlines guidelines and safeguards to ensure the collaboration remains focused -on its technical mission. - -The Policy applies to all Working Group activities and communications, -including official meetings, subcommittee discussions, emails, shared -documents, and any other interactions under the Working Group’s auspices (e.g., -group chats, version control systems). It applies to all participants from -Google, Anthropic, and any other member organizations or independent -contributors involved in the Working Group. Each participating entity should -ensure its representatives understand and uphold these rules. By participating, -members agree to follow the Policy in both letter and spirit. - -#### Governance Procedures and Principles - -- **Participant Guidelines.** Participants should generally be limited to - individuals in technical roles who are directly involved in the MCP SDK - project. These participants should not be key decision-makers in their - company’s AI commercial strategy, sales, marketing, pricing, or other - competitively or strategically sensitive business planning. -- **Agenda Preparation.** A written agenda should be circulated before each - Working Group meeting. Agenda items should focus on the SDK’s technical - development, maintenance, or documentation. Where appropriate, counsel should - review the agenda prior to circulation to ensure compliance with the Policy. -- **Policy Reminder at Start.** Meetings should begin with a brief antitrust - compliance reminder with reference to the Policy. -- **Minutes.** Meetings will be minuted by a designated participant, and - neutrally record attendees, roles, topics discussed, action items, and - outcomes. Draft minutes will be circulated to all participants. -- **Documentation and Transparency.** Working Group outputs are intended - for public release. Significant design decisions and discussion outcomes - should be documented publicly. If a topic cannot be safely disclosed - publicly, it likely does not belong in this forum. -- **Independence of Decision-Making.** All participants and their companies - retain complete independence in their own business decisions and competitive - strategies outside of the MCP SDK project. Nothing in this collaboration - restricts or influences how each company operates its commercial business. - -### Information Exchange Guidelines - -**Appropriate Topics.** Working Group members are anticipated to remain within the -project’s technical scope. In general, discussions should focus on improving -the Go SDK for MCP in a transparent, collaborative, and non-exclusive manner. -The following topics are appropriate and expected: - -- **Software Design and Architecture:** Implementation of MCP features in Go, - API design, code structure, testing frameworks, performance considerations, - compatibility issues, and security concerns. -- **Technical Contributions and Bug Fixes:** Review of code contributions, - debugging problems, and handling feature requests. -- **Documentation and Open-Source Logistics:** Discussions of project - documentation, changelogs, versioning strategy, and managing contributions. -- **Standards and Interoperability:** Ensuring SDK compliance with MCP or other - open technical standards. Any standardization effort should be open, - voluntary, and tailored to promote interoperability. -- **Public Information:** Any public information relevant to the Working - Group’s technical work (e.g., published research, open-source code from - outside projects, publicly documented API specs). - -**Inappropriate Topics.** To ensure compliance with antitrust law and maintain -the open character of the collaboration, the following subjects should not be -discussed in Working Group meetings, side conversations, or related communications: - -- **Pricing and Commercial Terms:** Do not discuss prices, pricing strategy, - discounting, or future pricing plans for Claude, Gemini, or any other AI - product or service provided by Working Group members. -- **Sales or Output:** Avoid sharing sales volumes, revenue, customer counts, - market shares, production plans, or any business performance metrics. -- **Product Roadmaps (beyond SDK):** Do not disclose internal plans for AI - model development, feature rollouts, or future commercialization strategies. -- **Customers or Markets:** No discussions of which customers, industries, or - geographies the parties will pursue or avoid. -- **Customer or Supplier Details:** Do not share specifics of contracts, - negotiations, or relationships with commercial partners. -- **Non-SDK Proprietary Tech:** Keep discussions focused on the open SDK. - Sharing of information should be limited to what is needed to achieve the - Working Group’s goals. Each party’s internal model architectures, fine-tuning - approaches, or training methods unrelated to the project should not be - disclosed. -- **HR or Labor Matters:** No discussions about wages, hiring plans, or - policies toward employees. - -#### Enforcement and Support - -- **Shared Responsibility.** All Working Group participants share responsibility - for upholding the Policy. While legal counsel can provide support, day-to-day - compliance is a function of culture and practice. -- **Designated Legal Contacts.** Each participating entity should designate a - legal point of contact responsible for reviewing meeting materials (e.g., - agendas, minutes) and fielding questions about compliance. These contacts - should be included in the Working Group distribution list for all official - materials and should be consulted in advance of any meetings where sensitive - topics may arise. -- **Final Note.** The Policy is not meant to chill legitimate technical - collaboration. It is meant to ensure that the Working Group can focus on its - mission without creating unnecessary legal risk or attracting regulator - scrutiny. Participants who follow the Policy and avoid Inappropriate Topics - will remain squarely within the procompetitive zone. - -### Ongoing evaluation - -On an ongoing basis, the administrators of the SDK will evaluate whether it is -keeping pace with changes to the MCP spec and meeting its goals of openness and -transparency. If it is not meeting these goals, either because it exceeds the -bandwidth of its current Working Group, or because the processes here are -inadequate, these processes will be re-evaluated by the Working Group. At this -time, the Working Group set may be expanded to include additional community -members, based on their history of strong contribution. +The goal of this repository is to provide a robust and complete Go +implementation of the model context protocol, in an open and transparent +manner, without bias toward specific integration paths or providers. To that +end, the model context protocol organization's [antitrust +policy](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/ANTITRUST.md) +applies to all participation in this project. diff --git a/LICENSE b/LICENSE index 508be92..5791499 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,193 @@ +The MCP project is undergoing a licensing transition from the MIT License to the Apache License, Version 2.0 ("Apache-2.0"). All new code and specification contributions to the project are licensed under Apache-2.0. Documentation contributions (excluding specifications) are licensed under CC-BY-4.0. + +Contributions for which relicensing consent has been obtained are licensed under Apache-2.0. Contributions made by authors who originally licensed their work under the MIT License and who have not yet granted explicit permission to relicense remain licensed under the MIT License. + +No rights beyond those granted by the applicable original license are conveyed for such contributions. + +--- + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright + owner or by an individual or Legal Entity authorized to submit on behalf + of the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + +--- + MIT License -Copyright (c) 2025 Go MCP SDK Authors +Copyright (c) 2024-2025 Model Context Protocol a Series of LF Projects, LLC. Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal @@ -19,3 +206,11 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +--- + +Creative Commons Attribution 4.0 International (CC-BY-4.0) + +Documentation in this project (excluding specifications) is licensed under +CC-BY-4.0. See https://creativecommons.org/licenses/by/4.0/legalcode for +the full license text. diff --git a/README.md b/README.md index 2a2d408..baeb3e7 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,7 @@ [![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/modelcontextprotocol/go-sdk) [![PkgGoDev](https://pkg.go.dev/badge/github.com/modelcontextprotocol/go-sdk)](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) +[![OpenSSF Scorecard](https://api.scorecard.dev/projects/github.com/modelcontextprotocol/go-sdk/badge)](https://scorecard.dev/viewer/?uri=github.com/modelcontextprotocol/go-sdk) This repository contains an implementation of the official Go software development kit (SDK) for the Model Context Protocol (MCP). @@ -29,6 +30,31 @@ The SDK consists of several importable packages: The SDK endeavors to implement the full MCP spec. The [`docs/`](/docs/) directory contains feature documentation, mapping the MCP spec to the packages above. +## Version Compatibility + +The following table shows which versions of the Go SDK support which versions of the MCP specification: + +| SDK Version | Latest MCP Spec | All Supported MCP Specs | +|-----------------|-------------------|------------------------------------------------------------------| +| v1.7.0+ | 2026-07-28 | 2026-07-28, 2025-11-25\*, 2025-06-18, 2025-03-26, 2024-11-05 | +| v1.4.0 - v1.6.1 | 2025-11-25\* | 2025-11-25\*, 2025-06-18, 2025-03-26, 2024-11-05 | +| v1.2.0 - v1.3.1 | 2025-11-25\*\* | 2025-11-25\*\*, 2025-06-18, 2025-03-26, 2024-11-05 | +| v1.0.0 - v1.1.0 | 2025-06-18 | 2025-06-18, 2025-03-26, 2024-11-05 | + +\* Client side OAuth has experimental support. + +\*\* Partial support for 2025-11-25 (client side OAuth and Sampling with tools not available). + +The roots, sampling, and logging features are deprecated as of protocol version +2026-07-28 by +[SEP-2577](https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging). +The SDK continues to support them for compatibility during the deprecation +window (at least twelve months). See the individual feature documentation for +migration guidance. + +New releases of the SDK target only supported versions of Go. See +https://go.dev/doc/devel/release#policy for more information. + ## Getting started To get started creating an MCP server, create an `mcp.Server` instance, add @@ -141,5 +167,5 @@ is a thriving ecosystem of Go MCP clients and servers. ## License -This project is licensed under the MIT License - see the [LICENSE](./LICENSE) -file for details. +This project is licensed under Apache 2.0 for new contributions, with existing +code under MIT - see the [LICENSE](./LICENSE) file for details. diff --git a/ROADMAP.md b/ROADMAP.md new file mode 100644 index 0000000..589bf55 --- /dev/null +++ b/ROADMAP.md @@ -0,0 +1,32 @@ +# MCP Go SDK Roadmap + +## Current focus + +The following items are planned for the next release (v1.4.0): + +- **SEP-1730: Tier 1 SDK support** (https://github.com/modelcontextprotocol/go-sdk/issues/675) + - Description: We aim to be rated as Tier 1 SDK. +- **SEP-1577: support Sampling With Tools** (https://github.com/modelcontextprotocol/go-sdk/issues/629) + - Description: Sampling with Tools is the last SEP to be implemented + to implement the full scope of 2025-11-25 MCP specification. +- **OAuth Support** (https://github.com/modelcontextprotocol/go-sdk/issues/19) + - Description: Client-side OAuth support is the last feature + needed to provide the Authorization solution as per specification. + +## Future work + +Once the tiering system is in place, we will draft a new roadmap to describe +our long-term strategy. It may include experimental features or extensions +to the MCP specification listed in the next section. + +## Experimental features + +- **SEP-1686 (experimental): Implement Tasks** (https://github.com/modelcontextprotocol/go-sdk/issues/626) + +## Extensions + +### Authorization extensions (https://github.com/modelcontextprotocol/ext-auth) + +- **ext-auth: Enterprise Managed Authorization** (https://github.com/modelcontextprotocol/go-sdk/issues/628) + - Status: In progress - https://github.com/modelcontextprotocol/go-sdk/pull/770 +- **ext-auth: OAuth Client Credentials** (https://github.com/modelcontextprotocol/go-sdk/issues/627) \ No newline at end of file diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 0000000..5029242 --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,21 @@ +# Security Policy + +Thank you for helping keep the Model Context Protocol and its ecosystem secure. + +## Reporting Security Issues + +If you discover a security vulnerability in this repository, please report it through +the [GitHub Security Advisory process](https://docs.github.com/en/code-security/security-advisories/guidance-on-reporting-and-writing-information-about-vulnerabilities/privately-reporting-a-security-vulnerability) +for this repository. + +Please **do not** report security vulnerabilities through public GitHub issues, discussions, +or pull requests. + +## What to Include + +To help us triage and respond quickly, please include: + +- A description of the vulnerability +- Steps to reproduce the issue +- The potential impact +- Any suggested fixes (optional) diff --git a/auth/auth.go b/auth/auth.go index 0eea1d8..dc02d7d 100644 --- a/auth/auth.go +++ b/auth/auth.go @@ -6,19 +6,27 @@ package auth import ( "context" + "encoding/json" "errors" + "fmt" "net/http" "slices" "strings" "time" + + "github.com/modelcontextprotocol/go-sdk/oauthex" ) // TokenInfo holds information from a bearer token. type TokenInfo struct { Scopes []string Expiration time.Time - // TODO: add standard JWT fields - Extra map[string]any + // UserID is an optional identifier for the authenticated user. + // If set by a TokenVerifier, it can be used by transports to prevent + // session hijacking by ensuring that all requests for a given session + // come from the same user. + UserID string + Extra map[string]any } // The error that a TokenVerifier should return if the token cannot be verified. @@ -39,6 +47,34 @@ type RequireBearerTokenOptions struct { ResourceMetadataURL string // The required scopes. Scopes []string + // AllowMissingExpiration opts the middleware out of the + // `tokenInfo.Expiration.IsZero()` reject. Default false preserves the + // existing strict behaviour (every TokenInfo must carry an Expiration). + // + // Some IdPs emit session-bound bearer tokens that do not carry a standalone + // `exp` claim — the token's lifetime is bounded by an external session and + // is not advertised in-band. Resource servers integrating with such IdPs + // need to opt in to validating the rest of the token (scopes, signature + // via the verifier callback, etc.) without requiring the expiration field + // to be present. + // + // When enabled, the verifier is still responsible for any session-level + // validity check it can perform; this option only relaxes the middleware's + // own expiration enforcement. + AllowMissingExpiration bool + // ClockSkew bounds the tolerance applied to a token's Expiration when + // deciding whether it has elapsed. A token is rejected only if + // Expiration + ClockSkew is before the current time. Zero (the default) + // preserves strict comparison: any expired token is rejected immediately. + // + // Resource servers running behind a CDN, in distributed deployments, or + // communicating with an authorization server whose clock drifts a few + // seconds (common with cloud-managed IdPs) need a small positive value + // here to avoid rejecting tokens that are valid by the issuer's clock + // but momentarily appear expired by the verifier's. The same tolerance + // guards against an issuer's clock running slightly fast at /token + // issuance time. + ClockSkew time.Duration } type tokenInfoKey struct{} @@ -66,8 +102,17 @@ func RequireBearerToken(verifier TokenVerifier, opts *RequireBearerTokenOptions) tokenInfo, errmsg, code := verify(r, verifier, opts) if code != 0 { if code == http.StatusUnauthorized || code == http.StatusForbidden { - if opts != nil && opts.ResourceMetadataURL != "" { - w.Header().Add("WWW-Authenticate", "Bearer resource_metadata="+opts.ResourceMetadataURL) + if opts != nil { + var params []string + if opts.ResourceMetadataURL != "" { + params = append(params, fmt.Sprintf("resource_metadata=%q", opts.ResourceMetadataURL)) + } + if len(opts.Scopes) > 0 { + params = append(params, fmt.Sprintf("scope=%q", strings.Join(opts.Scopes, " "))) + } + if len(params) > 0 { + w.Header().Add("WWW-Authenticate", "Bearer "+strings.Join(params, ", ")) + } } } http.Error(w, errmsg, code) @@ -98,6 +143,9 @@ func verify(req *http.Request, verifier TokenVerifier, opts *RequireBearerTokenO } return nil, err.Error(), http.StatusInternalServerError } + if tokenInfo == nil { + return nil, "token validation failed", http.StatusInternalServerError + } // Check scopes. All must be present. if opts != nil { @@ -109,12 +157,58 @@ func verify(req *http.Request, verifier TokenVerifier, opts *RequireBearerTokenO } } - // Check expiration. - if tokenInfo.Expiration.IsZero() { - return nil, "token missing expiration", http.StatusUnauthorized + if opts == nil { + opts = &RequireBearerTokenOptions{} } - if tokenInfo.Expiration.Before(time.Now()) { + // Check expiration, with optional clock-skew tolerance. Skew only applies + // when an expiration is present; a missing expiration is governed solely by + // AllowMissingExpiration. + if tokenInfo.Expiration.IsZero() { + if !opts.AllowMissingExpiration { + return nil, "token missing expiration", http.StatusUnauthorized + } + } else if tokenInfo.Expiration.Add(opts.ClockSkew).Before(time.Now()) { return nil, "token expired", http.StatusUnauthorized } return tokenInfo, "", 0 } + +// ProtectedResourceMetadataHandler returns an http.Handler that serves OAuth 2.0 +// protected resource metadata (RFC 9728) with CORS support. +// +// This handler allows cross-origin requests from any origin (Access-Control-Allow-Origin: *) +// because OAuth metadata is public information intended for client discovery (RFC 9728 §3.1). +// The metadata contains only non-sensitive configuration data about authorization servers +// and supported scopes. +// +// No validation of metadata fields is performed; ensure metadata accuracy at configuration time. +// +// For more sophisticated CORS policies or to restrict origins, wrap this handler with a +// CORS middleware like github.com/rs/cors or github.com/jub0bs/cors. +func ProtectedResourceMetadataHandler(metadata *oauthex.ProtectedResourceMetadata) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Set CORS headers for cross-origin client discovery. + // OAuth metadata is public information, so allowing any origin is safe. + w.Header().Set("Access-Control-Allow-Origin", "*") + w.Header().Set("Access-Control-Allow-Methods", "GET, OPTIONS") + w.Header().Set("Access-Control-Allow-Headers", "Content-Type") + + // Handle CORS preflight requests + if r.Method == http.MethodOptions { + w.WriteHeader(http.StatusNoContent) + return + } + + // Only GET allowed for metadata retrieval + if r.Method != http.MethodGet { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(metadata); err != nil { + http.Error(w, "Failed to encode metadata", http.StatusInternalServerError) + return + } + }) +} diff --git a/auth/auth_example_test.go b/auth/auth_example_test.go new file mode 100644 index 0000000..bf97eb8 --- /dev/null +++ b/auth/auth_example_test.go @@ -0,0 +1,222 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package auth_test + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "sync" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" +) + +// savingTokenSource is an oauth2.TokenSource that passes the oauth2 config and +// token to the given saver function each time the access token value changes. +type savingTokenSource struct { + mu sync.Mutex + src oauth2.TokenSource + saver func(*oauth2.Config, *oauth2.Token) error + config *oauth2.Config + accessToken string +} + +func (s *savingTokenSource) Token() (*oauth2.Token, error) { + s.mu.Lock() + defer s.mu.Unlock() + tok, err := s.src.Token() + if err != nil { + return nil, err + } + if s.accessToken != tok.AccessToken { + s.accessToken = tok.AccessToken + // This saver implementation always returns nil. + _ = s.saver(s.config, tok) + } + return tok, nil +} + +// NewSavingTokenSource persists OAuth 2.0 sessions by intercepting token +// changes from the wrapped oauth2.TokenSource. When this wrapper detects an +// access token change, it calls the provided session saver with the +// oauth2.Config and the new oauth2.Token. +// +// initial is an optional access token the caller may already hold (such as a +// token loaded from storage). It initialises the wrapper's state so that if +// the first wrapped.Token() call returns the same token, does not trigger a +// redundant call to saver(). Pass nil when there is no existing token, in +// which case the first token produced by wrapped.Token() is saved. +func NewSavingTokenSource(wrapped oauth2.TokenSource, config *oauth2.Config, initial *oauth2.Token, saver func(*oauth2.Config, *oauth2.Token) error) oauth2.TokenSource { + if wrapped == nil { + return nil + } + if saver == nil { + return wrapped + } + var accessToken string + if initial != nil { + accessToken = initial.AccessToken + } + return &savingTokenSource{ + src: wrapped, + saver: saver, + config: config, + accessToken: accessToken, + } +} + +// sessionStore is an in-memory OAuth 2.0 session store used to demonstrate +// persistence. A production implementation would persist the session to disk or +// a secret store. +type sessionStore struct { + config *oauth2.Config + token *oauth2.Token +} + +// save persists the given OAuth 2.0 config and token. +func (s *sessionStore) save(config *oauth2.Config, token *oauth2.Token) error { + fmt.Printf("Saving token: %s\n", token.AccessToken) + s.config = config + s.token = token + return nil +} + +// restore loads a previously persisted OAuth 2.0 session, if one exists. +func (s *sessionStore) restore() (*oauth2.Config, *oauth2.Token, error) { + if s.config != nil && s.token != nil { + fmt.Println("Restoring session.") + } else { + fmt.Println("No session found to restore.") + } + return s.config, s.token, nil +} + +// newMockAuthServer returns an httptest.Server that simulates both an MCP +// resource server requiring authorization and its OAuth authorization server. +func newMockAuthServer() *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/.well-known/oauth-authorization-server": + w.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprintf(w, `{"issuer": "http://%s", "authorization_endpoint": "http://%s/auth", "token_endpoint": "http://%s/token", "code_challenge_methods_supported": ["S256"]}`, r.Host, r.Host, r.Host) + case "/token": + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token": "mock-token", "token_type": "bearer"}`)) + default: + // The mock MCP endpoint returns 401 until the client presents a valid + // bearer token. + if r.Header.Get("Authorization") == "Bearer mock-token" { + // A real server would return a valid MCP message here. The empty + // response causes Connect to return an error after authorization, + // which is ignored for this example which demonstrates token + // persistence only. + w.WriteHeader(http.StatusOK) + return + } + w.Header().Set("WWW-Authenticate", "Bearer") + w.WriteHeader(http.StatusUnauthorized) + } + })) +} + +// This example shows how OAuth2 session persistence might be implemented. It +// connects twice using a shared in-memory session store: the first connection +// authorizes and saves the session, and the second restores and reuses it. +func Example_persistence() { + // Simulate an MCP server that requires authorization and an OAuth server. + mockServer := newMockAuthServer() + defer mockServer.Close() + + // store persists the OAuth2 session across both connections. + store := &sessionStore{} + + // connect performs a single client connection, restoring any saved session + // beforehand and saving any new session acquired during authorization. + connect := func() { + // Load the OAuth2 session if available, and use it. + var initialTS oauth2.TokenSource + cfg, tok, err := store.restore() + if err == nil && cfg != nil && tok != nil { + initialTS = NewSavingTokenSource( + cfg.TokenSource(context.Background(), tok), + cfg, tok, store.save, + ) + } + + // Configure and initialize the AuthorizationCodeHandler with a + // NewTokenSource that saves the session when the access token changes, + // and the InitialTokenSource set to the session loaded via restore(). + config := &auth.AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "example", + }, + Client: mockServer.Client(), + AuthorizationCodeFetcher: func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + fmt.Println("No token source found. Transport is calling Authorize()...") + // Extract the generated state from the authorization URL + u, _ := url.Parse(args.URL) + state := u.Query().Get("state") + return &auth.AuthorizationResult{Code: "mock-code", State: state}, nil + }, + NewTokenSource: func(ctx context.Context, cfg *oauth2.Config, token *oauth2.Token) (oauth2.TokenSource, error) { + // This save implementation always returns nil. + _ = store.save(cfg, token) + return NewSavingTokenSource( + cfg.TokenSource(ctx, token), cfg, token, store.save, + ), nil + }, + InitialTokenSource: initialTS, + } + handler, err := auth.NewAuthorizationCodeHandler(config) + if err != nil { + fmt.Printf("Error creating handler: %v\n", err) + return + } + + // Set the constructed handler on a transport. + transport := &mcp.StreamableClientTransport{ + Endpoint: mockServer.URL + "/sse", + OAuthHandler: handler, + } + + // Create a client and attempt to connect using the configured + // transport. The transport will automatically: + // 1. Call TokenSource() to check for an existing session. This will + // return InitialTokenSource, if set. + // 2. Try the MCP endpoint, and encounter a 401 response from the mock + // server. + // 3. Call Authorize() to perform the OAuth flow, which calls + // NewTokenSource. In this example NewTokenSource saves the newly + // acquired session. + // 4. Retry the MCP endpoint with a valid bearer token and get a 200. + client := mcp.NewClient( + &mcp.Implementation{Name: "example", Version: "1.0.0"}, nil, + ) + // Response ignored: this example asserts authorization only. + _, _ = client.Connect(context.Background(), transport, nil) + } + + // The first connection has no saved session, so it authorizes and saves. + fmt.Println("--- First connect ---") + connect() + // The second connection restores the saved session and reuses it, so no + // further authorization or save occurs. + fmt.Println("--- Second connect ---") + connect() + + // Output: + // --- First connect --- + // No session found to restore. + // No token source found. Transport is calling Authorize()... + // Saving token: mock-token + // --- Second connect --- + // Restoring session. +} diff --git a/auth/auth_test.go b/auth/auth_test.go index ef8ea7b..654ca26 100644 --- a/auth/auth_test.go +++ b/auth/auth_test.go @@ -6,10 +6,15 @@ package auth import ( "context" + "encoding/json" "errors" "net/http" + "net/http/httptest" + "strings" "testing" "time" + + "github.com/modelcontextprotocol/go-sdk/oauthex" ) func TestVerify(t *testing.T) { @@ -57,6 +62,11 @@ func TestVerify(t *testing.T) { "no expiration", nil, "Bearer noexp", "token missing expiration", 401, }, + { + "no expiration with AllowMissingExpiration accepts", + &RequireBearerTokenOptions{AllowMissingExpiration: true}, "Bearer noexp", + "", 0, + }, { "expired", nil, "Bearer expired", "token expired", 401, @@ -76,3 +86,266 @@ func TestVerify(t *testing.T) { }) } } + +func TestProtectedResourceMetadataHandler(t *testing.T) { + metadata := &oauthex.ProtectedResourceMetadata{ + Resource: "https://example.com/mcp", + AuthorizationServers: []string{ + "https://auth.example.com/.well-known/openid-configuration", + }, + ScopesSupported: []string{"read", "write"}, + } + + handler := ProtectedResourceMetadataHandler(metadata) + + tests := []struct { + name string + method string + wantStatus int + checkJSON bool + }{ + { + name: "GET returns metadata", + method: http.MethodGet, + wantStatus: http.StatusOK, + checkJSON: true, + }, + { + name: "OPTIONS for CORS preflight", + method: http.MethodOptions, + wantStatus: http.StatusNoContent, + }, + { + name: "POST not allowed", + method: http.MethodPost, + wantStatus: http.StatusMethodNotAllowed, + }, + { + name: "PUT not allowed", + method: http.MethodPut, + wantStatus: http.StatusMethodNotAllowed, + }, + { + name: "DELETE not allowed", + method: http.MethodDelete, + wantStatus: http.StatusMethodNotAllowed, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := httptest.NewRequest(tt.method, "/.well-known/oauth-protected-resource", nil) + rec := httptest.NewRecorder() + + handler.ServeHTTP(rec, req) + + if rec.Code != tt.wantStatus { + t.Errorf("status = %d, want %d", rec.Code, tt.wantStatus) + } + + // All responses should have CORS headers + if got := rec.Header().Get("Access-Control-Allow-Origin"); got != "*" { + t.Errorf("Access-Control-Allow-Origin = %q, want %q", got, "*") + } + + if got := rec.Header().Get("Access-Control-Allow-Methods"); got != "GET, OPTIONS" { + t.Errorf("Access-Control-Allow-Methods = %q, want %q", got, "GET, OPTIONS") + } + + // Validate error response body for disallowed methods + if tt.wantStatus == http.StatusMethodNotAllowed { + if !strings.Contains(rec.Body.String(), "Method not allowed") { + t.Errorf("error body = %q, want to contain %q", rec.Body.String(), "Method not allowed") + } + } + + if tt.checkJSON { + if got := rec.Header().Get("Content-Type"); got != "application/json" { + t.Errorf("Content-Type = %q, want %q", got, "application/json") + } + + var got oauthex.ProtectedResourceMetadata + if err := json.NewDecoder(rec.Body).Decode(&got); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + + if got.Resource != metadata.Resource { + t.Errorf("Resource = %q, want %q", got.Resource, metadata.Resource) + } + + if len(got.AuthorizationServers) != len(metadata.AuthorizationServers) { + t.Errorf("AuthorizationServers length = %d, want %d", + len(got.AuthorizationServers), len(metadata.AuthorizationServers)) + } + + for i, server := range got.AuthorizationServers { + if server != metadata.AuthorizationServers[i] { + t.Errorf("AuthorizationServers[%d] = %q, want %q", + i, server, metadata.AuthorizationServers[i]) + } + } + + if len(got.ScopesSupported) != len(metadata.ScopesSupported) { + t.Errorf("ScopesSupported length = %d, want %d", + len(got.ScopesSupported), len(metadata.ScopesSupported)) + } + } + }) + } +} + +func TestRequireBearerToken(t *testing.T) { + verifier := func(_ context.Context, token string, _ *http.Request) (*TokenInfo, error) { + if token == "valid" { + return &TokenInfo{Expiration: time.Now().Add(time.Hour), Scopes: []string{"read"}}, nil + } + return nil, ErrInvalidToken + } + + tests := []struct { + name string + opts *RequireBearerTokenOptions + authHeader string + wantHeader string + wantStatus int + }{ + { + name: "no middleware options", + opts: nil, + authHeader: "Bearer invalid", + wantHeader: "", + wantStatus: http.StatusUnauthorized, + }, + { + name: "metadata only", + opts: &RequireBearerTokenOptions{ + ResourceMetadataURL: "https://example.com/resource-metadata", + }, + authHeader: "Bearer invalid", + wantHeader: "Bearer resource_metadata=\"https://example.com/resource-metadata\"", + wantStatus: http.StatusUnauthorized, + }, + { + name: "scopes only", + opts: &RequireBearerTokenOptions{ + Scopes: []string{"read", "write"}, + }, + authHeader: "Bearer invalid", + wantHeader: "Bearer scope=\"read write\"", + wantStatus: http.StatusUnauthorized, + }, + { + name: "metadata and scopes", + opts: &RequireBearerTokenOptions{ + ResourceMetadataURL: "https://example.com/resource-metadata", + Scopes: []string{"read", "write"}, + }, + authHeader: "Bearer invalid", + wantHeader: "Bearer resource_metadata=\"https://example.com/resource-metadata\", scope=\"read write\"", + wantStatus: http.StatusUnauthorized, + }, + { + name: "insufficient scope", + opts: &RequireBearerTokenOptions{ + Scopes: []string{"admin"}, + }, + authHeader: "Bearer valid", // Has "read", needs "admin" -> 403 + wantHeader: "Bearer scope=\"admin\"", + wantStatus: http.StatusForbidden, + }, + { + name: "success", + opts: &RequireBearerTokenOptions{ + Scopes: []string{"read"}, + }, + authHeader: "Bearer valid", + wantHeader: "", + wantStatus: http.StatusOK, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + handler := RequireBearerToken(verifier, tt.opts)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest("GET", "/", nil) + if tt.authHeader != "" { + req.Header.Set("Authorization", tt.authHeader) + } + rec := httptest.NewRecorder() + + handler.ServeHTTP(rec, req) + + if rec.Code != tt.wantStatus { + t.Errorf("status = %d, want %d", rec.Code, tt.wantStatus) + } + + got := rec.Header().Get("WWW-Authenticate") + if got != tt.wantHeader { + t.Errorf("WWW-Authenticate = %q, want %q", got, tt.wantHeader) + } + }) + } +} + +// TestRequireBearerToken_ClockSkew verifies that the ClockSkew option +// extends the expiration check tolerance: a token whose Expiration is in the +// recent past is accepted iff the elapsed interval is within ClockSkew. +func TestRequireBearerToken_ClockSkew(t *testing.T) { + tests := []struct { + name string + clockSkew time.Duration + expiredAgo time.Duration + wantStatus int + }{ + { + name: "no skew, fresh token accepted", + clockSkew: 0, + expiredAgo: -time.Minute, // expires in 1 minute + wantStatus: http.StatusOK, + }, + { + name: "no skew, expired token rejected", + clockSkew: 0, + expiredAgo: 5 * time.Second, // expired 5s ago + wantStatus: http.StatusUnauthorized, + }, + { + name: "with skew, recently-expired token accepted", + clockSkew: 30 * time.Second, + expiredAgo: 5 * time.Second, + wantStatus: http.StatusOK, + }, + { + name: "with skew, token expired beyond tolerance rejected", + clockSkew: 10 * time.Second, + expiredAgo: 30 * time.Second, + wantStatus: http.StatusUnauthorized, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + verifier := func(_ context.Context, _ string, _ *http.Request) (*TokenInfo, error) { + return &TokenInfo{Expiration: time.Now().Add(-tt.expiredAgo)}, nil + } + handler := RequireBearerToken(verifier, &RequireBearerTokenOptions{ + ClockSkew: tt.clockSkew, + })(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest("GET", "/", nil) + req.Header.Set("Authorization", "Bearer anything") + rec := httptest.NewRecorder() + + handler.ServeHTTP(rec, req) + + if rec.Code != tt.wantStatus { + t.Errorf("status = %d, want %d", rec.Code, tt.wantStatus) + } + }) + } +} diff --git a/auth/authorization_code.go b/auth/authorization_code.go new file mode 100644 index 0000000..f349842 --- /dev/null +++ b/auth/authorization_code.go @@ -0,0 +1,713 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package auth + +import ( + "context" + "crypto/rand" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "slices" + "strings" + "sync" + + "github.com/modelcontextprotocol/go-sdk/internal/authutil" + "github.com/modelcontextprotocol/go-sdk/internal/util" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" +) + +// ClientIDMetadataDocumentConfig is used to configure the Client ID Metadata Document +// based client registration per +// https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#client-id-metadata-documents. +// See https://client.dev/ for more information. +type ClientIDMetadataDocumentConfig struct { + // URL is the client identifier URL as per + // https://datatracker.ietf.org/doc/html/draft-ietf-oauth-client-id-metadata-document-00#section-3. + URL string +} + +// DynamicClientRegistrationConfig is used to configure dynamic client registration per +// https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#dynamic-client-registration. +type DynamicClientRegistrationConfig struct { + // Metadata to be used in dynamic client registration request as per + // https://datatracker.ietf.org/doc/html/rfc7591#section-2. + // + // If Metadata.ApplicationType is empty, it will be inferred from + // Metadata.RedirectURIs. When set, it will be validated against the inferred type + // and an error will be returned if they conflict. + Metadata *oauthex.ClientRegistrationMetadata +} + +// AuthorizationResult is the result of an authorization flow. +// It is returned by [AuthorizationCodeHandler].AuthorizationCodeFetcher implementations. +type AuthorizationResult struct { + // Code is the authorization code obtained from the authorization server. + Code string + // State string returned by the authorization server. + State string + // Iss is the issuer identifier returned by the authorization server in the + // authorization response per [RFC 9207]. The AuthorizationCodeFetcher should + // populate this from the "iss" query parameter in the redirect URI if present. + // + // [RFC 9207]: https://www.rfc-editor.org/rfc/rfc9207 + Iss string +} + +// AuthorizationArgs is the input to [AuthorizationCodeFetcher]. +type AuthorizationArgs struct { + // Authorization URL to be opened in a browser for the user to start the authorization process. + URL string +} + +// AuthorizationCodeFetcher is called to initiate the OAuth authorization flow. +// It is responsible for directing the user to the authorization URL (e.g., opening +// in a browser) and returning the authorization code and state once the Authorization +// Server redirects back to the configured RedirectURL. +type AuthorizationCodeFetcher func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) + +// AuthorizationCodeHandlerConfig is the configuration for [AuthorizationCodeHandler]. +type AuthorizationCodeHandlerConfig struct { + // Client registration configuration. + // It is attempted in the following order: + // 1. Client ID Metadata Document + // 2. Preregistration + // 3. Dynamic Client Registration + // At least one method must be configured. + ClientIDMetadataDocumentConfig *ClientIDMetadataDocumentConfig + PreregisteredClient *oauthex.ClientCredentials + DynamicClientRegistrationConfig *DynamicClientRegistrationConfig + + // RedirectURL is a required URL to redirect to after authorization. + // The caller is responsible for handling the redirect out of band. + // + // If Dynamic Client Registration is used: + // - this field is permitted to be empty, in which case it will be set + // to the first redirect URI from + // DynamicClientRegistrationConfig.Metadata.RedirectURIs. + // - if the field is not empty, it must be one of the redirect URIs in + // DynamicClientRegistrationConfig.Metadata.RedirectURIs. + RedirectURL string + + // AuthorizationCodeFetcher is a required function called to initiate the authorization flow. + // See [AuthorizationCodeFetcher] for details. + AuthorizationCodeFetcher AuthorizationCodeFetcher + + // ScopeFilter, if non-nil, is called with the scopes discovered from the + // protected resource metadata (or WWW-Authenticate challenge) and returns + // the scopes to request during authorization. It gives the client full + // control to narrow, extend, or reorder that set — e.g. dropping Gmail's + // gmail.metadata, which the Gmail API refuses to combine with the search "q" + // parameter even alongside gmail.readonly. It runs before offline_access + // (see RequestRefreshToken) and the step-up union, so neither is affected. + ScopeFilter func(discovered []string) []string + + // RequestRefreshToken indicates that the client intends to use refresh + // tokens and is capable of storing them securely. + // + // When true and the Authorization Server metadata contains "offline_access" + // in its scopes_supported, the client adds "offline_access" to the + // requested scopes. + // + // When using Dynamic Client Registration, callers should include + // "refresh_token" in [DynamicClientRegistrationConfig].Metadata.GrantTypes + // directly to advertise refresh token support to the Authorization Server. + // + // When using Client ID Metadata Document, the document hosted at the + // Client ID URL should include "refresh_token" in its grant_types. + // + // See https://modelcontextprotocol.io/seps/2207-oidc-refresh-token-guidance. + RequestRefreshToken bool + + // AcceptUnadvertisedIss accepts a matching RFC 9207 iss even when the + // authorization server metadata omits authorization_response_iss_parameter_supported. + // The zero value (false) keeps the historical reject-unadvertised-iss behavior. + AcceptUnadvertisedIss bool + + // Client is an optional HTTP client to use for HTTP requests. + // It is used for the following requests: + // - Fetching Protected Resource Metadata + // - Fetching Authorization Server Metadata + // - Registering a client dynamically + // - Exchanging an authorization code for an access token + // - Refreshing an access token + // Custom clients can include additional security configurations, + // such as SSRF protections, see + // https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#server-side-request-forgery-ssrf + // If not provided, http.DefaultClient will be used. + Client *http.Client + + // NewTokenSource is an optional function that can be set to construct the + // token source that will be used by the [AuthorizationCodeHandler]. If + // non-nil, it is called after the authorization code is successfully + // exchanged for a token in [AuthorizationCodeHandler.Authorize] + // to obtain the [oauth2.TokenSource] returned by + // [AuthorizationCodeHandler.TokenSource]. Implementations must use the + // provided context, which is properly configured for constructing a + // TokenSource. The default is to call [oauth2.Config.TokenSource]. + NewTokenSource func(context.Context, *oauth2.Config, *oauth2.Token) (oauth2.TokenSource, error) + + // InitialTokenSource is an optional field that can be set to inject the + // token source that will be used by the [AuthorizationCodeHandler]. If + // non-nil, it is set as the token source that will be returned by + // [AuthorizationCodeHandler.TokenSource] during handler initialization. + // The default is nil, which means no token source has been set initially, + // and will trigger a call to [AuthorizationCodeHandler.Authorize]. + InitialTokenSource oauth2.TokenSource +} + +// AuthorizationCodeHandler is an implementation of [OAuthHandler] that uses +// the authorization code flow to obtain access tokens. +type AuthorizationCodeHandler struct { + config *AuthorizationCodeHandlerConfig + + // mu protects concurrent access to tokenSource and grantedScopes. + mu sync.RWMutex + + // tokenSource is the token source to use for authorization. + tokenSource oauth2.TokenSource + + // grantedScopes maps authorization server issuer to the list of scopes granted by that issuer. + grantedScopes map[string][]string +} + +var _ OAuthHandler = (*AuthorizationCodeHandler)(nil) + +func (h *AuthorizationCodeHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) { + h.mu.RLock() + defer h.mu.RUnlock() + return h.tokenSource, nil +} + +// NewAuthorizationCodeHandler creates a new AuthorizationCodeHandler. +// It performs validation of the configuration and returns an error if it is invalid. +// The passed config is consumed by the handler and should not be modified after. +func NewAuthorizationCodeHandler(config *AuthorizationCodeHandlerConfig) (*AuthorizationCodeHandler, error) { + if config == nil { + return nil, errors.New("config must be provided") + } + if config.ClientIDMetadataDocumentConfig == nil && + config.PreregisteredClient == nil && + config.DynamicClientRegistrationConfig == nil { + return nil, errors.New("at least one client registration configuration must be provided") + } + if config.AuthorizationCodeFetcher == nil { + return nil, errors.New("AuthorizationCodeFetcher is required") + } + if config.ClientIDMetadataDocumentConfig != nil && !isNonRootHTTPSURL(config.ClientIDMetadataDocumentConfig.URL) { + return nil, fmt.Errorf("client ID metadata document URL must be a non-root HTTPS URL") + } + if config.PreregisteredClient != nil { + if err := config.PreregisteredClient.Validate(); err != nil { + return nil, fmt.Errorf("invalid PreregisteredClient configuration: %w", err) + } + } + dCfg := config.DynamicClientRegistrationConfig + if dCfg != nil { + if dCfg.Metadata == nil { + return nil, errors.New("dynamic client registration requires non-nil Metadata") + } + if len(dCfg.Metadata.RedirectURIs) == 0 { + return nil, errors.New("Metadata.RedirectURIs is required for dynamic client registration") + } + if config.RedirectURL == "" { + config.RedirectURL = dCfg.Metadata.RedirectURIs[0] + } else if !slices.Contains(dCfg.Metadata.RedirectURIs, config.RedirectURL) { + return nil, fmt.Errorf("RedirectURL %q is not in the list of allowed redirect URIs for dynamic client registration", config.RedirectURL) + } + applicationType := inferApplicationType(dCfg.Metadata.RedirectURIs) + if dCfg.Metadata.ApplicationType == "" { + dCfg.Metadata.ApplicationType = applicationType + } else if dCfg.Metadata.ApplicationType != applicationType { + return nil, fmt.Errorf("application type %q conflicts with the application type inferred from redirect URIs", dCfg.Metadata.ApplicationType) + } + } + if config.RedirectURL == "" { + // If the RedirectURL was supposed to be set by the dynamic client registration, + // it should have been set by now. Otherwise, it is required. + return nil, errors.New("RedirectURL is required") + } + if config.Client == nil { + config.Client = http.DefaultClient + } + return &AuthorizationCodeHandler{ + config: config, + tokenSource: config.InitialTokenSource, + grantedScopes: make(map[string][]string), + }, nil +} + +func isNonRootHTTPSURL(u string) bool { + pu, err := url.Parse(u) + if err != nil { + return false + } + return pu.Scheme == "https" && pu.Path != "" +} + +// inferApplicationType returns an application type based on the redirect URIs. +func inferApplicationType(redirectURIs []string) string { + hasNative := false + hasWeb := false + for _, uri := range redirectURIs { + u, err := url.Parse(uri) + if err != nil { + return "" + } + switch u.Scheme { + case "http", "https": + if util.IsLoopback(u.Hostname()) { + hasNative = true + } else { + hasWeb = true + } + default: + hasNative = true + } + } + + if hasNative && hasWeb { + return "" + } + if hasNative { + return "native" + } + return "web" +} + +// Authorize performs the authorization flow. +// It is designed to perform the whole Authorization Code Grant flow. +// On success, [AuthorizationCodeHandler.TokenSource] will return a token source with the fetched token. +func (h *AuthorizationCodeHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error { + defer resp.Body.Close() + defer io.Copy(io.Discard, resp.Body) + + wwwChallenges, err := oauthex.ParseWWWAuthenticate(resp.Header[http.CanonicalHeaderKey("WWW-Authenticate")]) + if err != nil { + return fmt.Errorf("failed to parse WWW-Authenticate header: %v", err) + } + + if resp.StatusCode == http.StatusForbidden && errorFromChallenges(wwwChallenges) != "insufficient_scope" { + // We only want to perform step-up authorization for insufficient_scope errors. + // Returning nil, so that the call is retried immediately and the response + // is handled appropriately by the connection. + // Step-up authorization is defined at + // https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#step-up-authorization-flow + return nil + } + + prm, err := h.getProtectedResourceMetadata(ctx, wwwChallenges, req.URL.String()) + if err != nil { + return err + } + + asm, err := GetAuthServerMetadata(ctx, prm.AuthorizationServers[0], h.config.Client) + if err != nil { + return fmt.Errorf("failed to get authorization server metadata: %w", err) + } + if asm == nil { + // Fallback to 2025-03-26 spec: predefined endpoints. + // https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization#fallbacks-for-servers-without-metadata-discovery + authServerURL := prm.AuthorizationServers[0] + asm = &oauthex.AuthServerMeta{ + Issuer: authServerURL, + AuthorizationEndpoint: authServerURL + "/authorize", + TokenEndpoint: authServerURL + "/token", + RegistrationEndpoint: authServerURL + "/register", + } + } + + resolvedClientConfig, err := h.handleRegistration(ctx, asm) + if err != nil { + return err + } + + requestedScopes := scopesFromChallenges(wwwChallenges) + if len(requestedScopes) == 0 && len(prm.ScopesSupported) > 0 { + requestedScopes = prm.ScopesSupported + } + + // Let the client adjust the discovered scopes before offline_access and the + // step-up union below so neither is affected. + if h.config.ScopeFilter != nil { + requestedScopes = h.config.ScopeFilter(requestedScopes) + } + + // SEP-2207: when the client desires refresh tokens and the Authorization + // Server advertises offline_access support, add it to the requested scopes. + if h.config.RequestRefreshToken && + slices.Contains(asm.ScopesSupported, "offline_access") && + !slices.Contains(requestedScopes, "offline_access") { + requestedScopes = append(requestedScopes, "offline_access") + } + + // Accumulate scopes: union previously granted scopes with the newly + // challenged scopes so that step-up authorization does not lose + // permissions granted in earlier rounds (SEP-2350). + h.mu.RLock() + granted := h.grantedScopes[asm.Issuer] + h.mu.RUnlock() + requestedScopes = authutil.UnionScopes(granted, requestedScopes) + + cfg := &oauth2.Config{ + ClientID: resolvedClientConfig.clientID, + ClientSecret: resolvedClientConfig.clientSecret, + + Endpoint: oauth2.Endpoint{ + AuthURL: asm.AuthorizationEndpoint, + TokenURL: asm.TokenEndpoint, + AuthStyle: resolvedClientConfig.authStyle, + }, + RedirectURL: h.config.RedirectURL, + Scopes: requestedScopes, + } + + authRes, err := h.getAuthorizationCode(ctx, cfg, prm.Resource) + if err != nil { + // Purposefully leaving the error unwrappable so it can be handled by the caller. + return err + } + if err := validateIssuerResponse(authRes.Iss, asm.Issuer, asm.AuthorizationResponseIssParameterSupported, h.config.AcceptUnadvertisedIss); err != nil { + return err + } + + err = h.exchangeAuthorizationCode(ctx, cfg, authRes, prm.Resource) + if err != nil { + return err + } + + return h.updateGrantedScopes(asm.Issuer, requestedScopes) +} + +// resourceMetadataURLFromChallenges returns a resource metadata URL from the given "WWW-Authenticate" header challenges, +// or the empty string if there is none. +func resourceMetadataURLFromChallenges(cs []oauthex.Challenge) string { + for _, c := range cs { + if u := c.Params["resource_metadata"]; u != "" { + return u + } + } + return "" +} + +// scopesFromChallenges returns the scopes from the given "WWW-Authenticate" header challenges. +// It only looks at challenges with the "Bearer" scheme. +func scopesFromChallenges(cs []oauthex.Challenge) []string { + for _, c := range cs { + if c.Scheme == "bearer" && c.Params["scope"] != "" { + return strings.Fields(c.Params["scope"]) + } + } + return nil +} + +// errorFromChallenges returns the error from the given "WWW-Authenticate" header challenges. +// It only looks at challenges with the "Bearer" scheme. +func errorFromChallenges(cs []oauthex.Challenge) string { + for _, c := range cs { + if c.Scheme == "bearer" && c.Params["error"] != "" { + return c.Params["error"] + } + } + return "" +} + +// getProtectedResourceMetadata returns the protected resource metadata. +// If no metadata was found or the fetched metadata fails security checks, +// it returns an error. +func (h *AuthorizationCodeHandler) getProtectedResourceMetadata(ctx context.Context, wwwChallenges []oauthex.Challenge, mcpServerURL string) (*oauthex.ProtectedResourceMetadata, error) { + // Use MCP server URL as the resource URI per + // https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#canonical-server-uri. + for _, url := range protectedResourceMetadataURLs(resourceMetadataURLFromChallenges(wwwChallenges), mcpServerURL) { + prm, err := oauthex.GetProtectedResourceMetadata(ctx, url.URL, url.Resource, h.config.Client) + if err != nil { + continue + } + if prm == nil { + continue + } + if len(prm.AuthorizationServers) == 0 { + // If we found PRM, we enforce the 2025-11-25 spec and not search further. + return nil, fmt.Errorf("protected resource metadata has no authorization servers specified") + } + return prm, nil + } + // Fallback to 2025-03-26 spec MCP server root is the Authorization Server: + // https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization#server-metadata-discovery + u, err := url.Parse(mcpServerURL) + if err != nil { + return nil, fmt.Errorf("failed to parse MCP server URL: %v", err) + } + u.Path = "" + prm := &oauthex.ProtectedResourceMetadata{ + AuthorizationServers: []string{u.String()}, + Resource: mcpServerURL, + } + return prm, nil +} + +type prmURL struct { + // URL represents a URL where Protected Resource Metadata may be retrieved. + URL string + // Resource represents the corresponding resource URL for [URL]. + // It is required to perform validation described in RFC 9728, section 3.3. + Resource string +} + +// protectedResourceMetadataURLs returns a list of URLs to try when looking for +// protected resource metadata as mandated by the MCP specification: +// https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#protected-resource-metadata-discovery-requirements +func protectedResourceMetadataURLs(metadataURL, resourceURL string) []prmURL { + var urls []prmURL + if metadataURL != "" { + urls = append(urls, prmURL{ + URL: metadataURL, + Resource: resourceURL, + }) + } + ru, err := url.Parse(resourceURL) + if err != nil { + return urls + } + mu := *ru + // "At the path of the server's MCP endpoint". + mu.Path = "/.well-known/oauth-protected-resource/" + strings.TrimLeft(ru.Path, "/") + urls = append(urls, prmURL{ + URL: mu.String(), + Resource: resourceURL, + }) + // "At the root". + mu.Path = "/.well-known/oauth-protected-resource" + ru.Path = "" + urls = append(urls, prmURL{ + URL: mu.String(), + Resource: ru.String(), + }) + return urls +} + +type registrationType int + +const ( + registrationTypeClientIDMetadataDocument registrationType = iota + registrationTypePreregistered + registrationTypeDynamic +) + +type resolvedClientConfig struct { + registrationType registrationType + clientID string + clientSecret string + authStyle oauth2.AuthStyle +} + +func selectTokenAuthMethod(supported []string) oauth2.AuthStyle { + prefOrder := []string{ + // Preferred in OAuth 2.1 draft: https://www.ietf.org/archive/id/draft-ietf-oauth-v2-1-14.html#name-client-secret. + "client_secret_post", + "client_secret_basic", + } + for _, method := range prefOrder { + if slices.Contains(supported, method) { + return authMethodToStyle(method) + } + } + return oauth2.AuthStyleAutoDetect +} + +func authMethodToStyle(method string) oauth2.AuthStyle { + switch method { + case "client_secret_post": + return oauth2.AuthStyleInParams + case "client_secret_basic": + return oauth2.AuthStyleInHeader + case "none": + // "none" is equivalent to "client_secret_post" but without sending client secret. + return oauth2.AuthStyleInParams + default: + // "client_secret_basic" is the default per https://datatracker.ietf.org/doc/html/rfc7591#section-2. + return oauth2.AuthStyleInHeader + } +} + +// handleRegistration handles client registration. +// The provided authorization server metadata must be non-nil. +// Support for different registration methods is defined as follows: +// - Client ID Metadata Document: metadata must have +// `ClientIDMetadataDocumentSupported` set to true. +// - Pre-registered client: assumed to be supported. +// - Dynamic client registration: metadata must have +// `RegistrationEndpoint` set to a non-empty value. +func (h *AuthorizationCodeHandler) handleRegistration(ctx context.Context, asm *oauthex.AuthServerMeta) (*resolvedClientConfig, error) { + // 1. Attempt to use Client ID Metadata Document (SEP-991). + cimdCfg := h.config.ClientIDMetadataDocumentConfig + if cimdCfg != nil && asm.ClientIDMetadataDocumentSupported { + return &resolvedClientConfig{ + registrationType: registrationTypeClientIDMetadataDocument, + clientID: cimdCfg.URL, + }, nil + } + // 2. Attempt to use pre-registered client configuration. + preCfg := h.config.PreregisteredClient + if preCfg != nil { + if preCfg.Issuer != "" && !authutil.IssuersEqual(preCfg.Issuer, asm.Issuer) { + return nil, fmt.Errorf("authorization server issuer %q does not match pre-registered credentials issuer %q", asm.Issuer, preCfg.Issuer) + } + authStyle := selectTokenAuthMethod(asm.TokenEndpointAuthMethodsSupported) + clientSecret := "" + if preCfg.ClientSecretAuth != nil { + clientSecret = preCfg.ClientSecretAuth.ClientSecret + } + return &resolvedClientConfig{ + registrationType: registrationTypePreregistered, + clientID: preCfg.ClientID, + clientSecret: clientSecret, + authStyle: authStyle, + }, nil + } + // 3. Attempt to use dynamic client registration. + dcrCfg := h.config.DynamicClientRegistrationConfig + if dcrCfg != nil && asm.RegistrationEndpoint != "" { + regResp, err := oauthex.RegisterClient(ctx, asm.RegistrationEndpoint, dcrCfg.Metadata, h.config.Client) + if err != nil { + return nil, fmt.Errorf("failed to register client: %w", err) + } + cfg := &resolvedClientConfig{ + registrationType: registrationTypeDynamic, + clientID: regResp.ClientID, + clientSecret: regResp.ClientSecret, + authStyle: authMethodToStyle(regResp.TokenEndpointAuthMethod), + } + return cfg, nil + } + return nil, fmt.Errorf("no configured client registration methods are supported by the authorization server") +} + +type authResult struct { + *AuthorizationResult + // usedCodeVerifier is the PKCE code verifier used to obtain the authorization code. + // It is preserved for the token exchange step. + usedCodeVerifier string +} + +// getAuthorizationCode uses the [AuthorizationCodeHandler.AuthorizationCodeFetcher] +// to obtain an authorization code. +func (h *AuthorizationCodeHandler) getAuthorizationCode(ctx context.Context, cfg *oauth2.Config, resourceURL string) (*authResult, error) { + codeVerifier := oauth2.GenerateVerifier() + state := rand.Text() + + authURL := cfg.AuthCodeURL(state, + oauth2.S256ChallengeOption(codeVerifier), + oauth2.SetAuthURLParam("resource", resourceURL), + ) + + authRes, err := h.config.AuthorizationCodeFetcher(ctx, &AuthorizationArgs{URL: authURL}) + if err != nil { + // Purposefully leaving the error unwrappable so it can be handled by the caller. + return nil, err + } + if authRes.State != state { + return nil, fmt.Errorf("state mismatch") + } + return &authResult{ + AuthorizationResult: authRes, + usedCodeVerifier: codeVerifier, + }, nil +} + +// validateIssuerResponse validates the "iss" parameter in an authorization response +// per [RFC 9207]. When the server advertises authorization_response_iss_parameter_supported, +// iss is required and must match expectedIssuer. When it does not advertise support, +// an empty iss is accepted; a present iss is compared to expectedIssuer (RFC 9207 §2.4) +// and a mismatch is always rejected. A matching unadvertised iss is accepted only when +// acceptUnadvertisedIss is true (local policy); the default is the historical reject. +// +// [RFC 9207]: https://www.rfc-editor.org/rfc/rfc9207 +func validateIssuerResponse(iss, expectedIssuer string, issParameterSupported, acceptUnadvertisedIss bool) error { + if issParameterSupported { + if iss == "" { + return fmt.Errorf("authorization server advertises RFC 9207 iss parameter support but none was received in the authorization response") + } + if iss != expectedIssuer { + return fmt.Errorf("authorization response issuer %q does not match expected issuer %q", iss, expectedIssuer) + } + return nil + } + if iss == "" { + return nil + } + if iss != expectedIssuer { + return fmt.Errorf("authorization response issuer %q does not match expected issuer %q", iss, expectedIssuer) + } + if !acceptUnadvertisedIss { + return fmt.Errorf("authorization server does not advertise RFC 9207 iss parameter support but iss was received in the authorization response") + } + return nil +} + +// exchangeAuthorizationCode exchanges the authorization code for a token +// and stores it in a token source. +func (h *AuthorizationCodeHandler) exchangeAuthorizationCode(ctx context.Context, cfg *oauth2.Config, authResult *authResult, resourceURL string) error { + opts := []oauth2.AuthCodeOption{ + oauth2.VerifierOption(authResult.usedCodeVerifier), + oauth2.SetAuthURLParam("resource", resourceURL), + } + clientCtx := context.WithValue(ctx, oauth2.HTTPClient, h.config.Client) + token, err := cfg.Exchange(clientCtx, authResult.Code, opts...) + if err != nil { + return fmt.Errorf("token exchange failed: %w", err) + } + // The token source outlives this authorization request: it is stored on the + // handler and used by the transport for the lifetime of the connection. The + // oauth2 library captures the context passed to TokenSource and reuses it for + // every subsequent token refresh (see golang.org/x/oauth2: tokenRefresher + // retains the context and passes it to each refresh round-trip). Binding it to + // the per-request ctx makes all later refreshes fail with "context canceled" + // once that request (or the connect operation that triggered authorization) + // completes. Use a background context that still carries the configured HTTP + // client so refreshes keep working for the life of the token source. + refreshCtx := context.WithValue(context.Background(), oauth2.HTTPClient, h.config.Client) + var ts oauth2.TokenSource + if h.config.NewTokenSource == nil { + ts = cfg.TokenSource(refreshCtx, token) + } else { + var err error + ts, err = h.config.NewTokenSource(refreshCtx, cfg, token) + if err != nil { + return fmt.Errorf("constructing token source failed: %w", err) + } + } + h.mu.Lock() + h.tokenSource = ts + h.mu.Unlock() + return nil +} + +// updateGrantedScopes updates the granted scopes based on the token source and requested scopes. +func (h *AuthorizationCodeHandler) updateGrantedScopes(issuer string, requestedScopes []string) error { + h.mu.RLock() + ts := h.tokenSource + h.mu.RUnlock() + + if ts == nil { + return nil + } + tok, err := ts.Token() + if err != nil { + return err + } + + h.mu.Lock() + if tokenScopes := authutil.ScopesFromToken(tok); tokenScopes == nil { + h.grantedScopes[issuer] = requestedScopes + } else { + h.grantedScopes[issuer] = tokenScopes + } + h.mu.Unlock() + return nil +} diff --git a/auth/authorization_code_test.go b/auth/authorization_code_test.go new file mode 100644 index 0000000..69a7ca6 --- /dev/null +++ b/auth/authorization_code_test.go @@ -0,0 +1,1562 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package auth + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "net/http/httputil" + "net/url" + "slices" + "strings" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" + "github.com/modelcontextprotocol/go-sdk/internal/oauthtest" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" +) + +func TestAuthorize(t *testing.T) { + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test_client_id": { + Secret: "test_client_secret", + RedirectURIs: []string{"http://localhost:12345/callback"}, + }, + }, + }, + }) + authServer.Start(t) + + resourceMux := http.NewServeMux() + resourceServer := httptest.NewServer(resourceMux) + t.Cleanup(resourceServer.Close) + resourceURL := resourceServer.URL + "/resource" + + resourceMux.Handle("/.well-known/oauth-protected-resource/resource", ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: resourceURL, + AuthorizationServers: []string{authServer.URL()}, + })) + + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost:12345/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "test_client_secret", + }, + }, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + // The fake authorization server will redirect to an URL with code and state. + client := &http.Client{ + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } + resp, err := client.Get(args.URL) + if err != nil { + return nil, fmt.Errorf("failed to visit auth URL: %v", err) + } + defer resp.Body.Close() + dump, err := httputil.DumpResponse(resp, true) + if err != nil { + t.Fatalf("failed to dump response: %v", err) + } + t.Log(string(dump)) + + location, err := resp.Location() + if err != nil { + return nil, fmt.Errorf("failed to get location header: %v", err) + } + return &AuthorizationResult{ + Code: location.Query().Get("code"), + State: location.Query().Get("state"), + Iss: location.Query().Get("iss"), + }, nil + }, + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, resourceURL, nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + resp.Header.Set( + "WWW-Authenticate", + "Bearer resource_metadata="+resourceServer.URL+"/.well-known/oauth-protected-resource/resource", + ) + + if err := handler.Authorize(context.Background(), req, resp); err != nil { + t.Fatalf("Authorize failed: %v", err) + } + + tokenSource, err := handler.TokenSource(t.Context()) + if err != nil { + t.Fatalf("Failed to get token source: %v", err) + } + token, err := tokenSource.Token() + if err != nil { + t.Fatalf("Failed to get token: %v", err) + } + if token.AccessToken != "test_access_token" { + t.Errorf("Expected access token 'test_access_token', got '%s'", token.AccessToken) + } +} + +// TestAuthorize_RefreshAfterContextCancel verifies that the token source built +// by Authorize keeps refreshing after the context passed to Authorize is +// cancelled. The token source is stored on the handler and used by the +// transport for the whole connection lifetime, whereas Authorize is typically +// called from a request- or connect-scoped context. golang.org/x/oauth2 +// captures the context handed to Config.TokenSource and reuses it for every +// refresh, so binding it to the request context made every refresh after the +// access token expired fail with "context canceled". Regression test for that. +func TestAuthorize_RefreshAfterContextCancel(t *testing.T) { + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + // expires_in below oauth2's 10s expiry delta, so the reuse token source + // treats the access token as expired immediately and must refresh. + AccessTokenTTL: 1, + IssueRefreshToken: true, + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test_client_id": { + Secret: "test_client_secret", + RedirectURIs: []string{"http://localhost:12345/callback"}, + }, + }, + }, + }) + authServer.Start(t) + + resourceMux := http.NewServeMux() + resourceServer := httptest.NewServer(resourceMux) + t.Cleanup(resourceServer.Close) + resourceURL := resourceServer.URL + "/resource" + resourceMux.Handle("/.well-known/oauth-protected-resource/resource", ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: resourceURL, + AuthorizationServers: []string{authServer.URL()}, + })) + + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost:12345/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ClientSecret: "test_client_secret"}, + }, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + client := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + resp, err := client.Get(args.URL) + if err != nil { + return nil, fmt.Errorf("failed to visit auth URL: %v", err) + } + defer resp.Body.Close() + location, err := resp.Location() + if err != nil { + return nil, fmt.Errorf("failed to get location header: %v", err) + } + return &AuthorizationResult{ + Code: location.Query().Get("code"), + State: location.Query().Get("state"), + Iss: location.Query().Get("iss"), + }, nil + }, + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, resourceURL, nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + resp.Header.Set("WWW-Authenticate", "Bearer resource_metadata="+resourceServer.URL+"/.well-known/oauth-protected-resource/resource") + + // Authorize under a context that we cancel immediately afterwards, mimicking + // the request/connect context the transport passes and that is already done + // by the time a later token refresh runs. + ctx, cancel := context.WithCancel(context.Background()) + if err := handler.Authorize(ctx, req, resp); err != nil { + t.Fatalf("Authorize failed: %v", err) + } + cancel() + + tokenSource, err := handler.TokenSource(context.Background()) + if err != nil { + t.Fatalf("Failed to get token source: %v", err) + } + // The access token is already expired, so this forces a refresh round-trip. + // It must not fail with "context canceled" from the cancelled Authorize ctx. + token, err := tokenSource.Token() + if err != nil { + t.Fatalf("token refresh after Authorize context cancellation failed: %v", err) + } + if token.AccessToken != "test_access_token_refreshed" { + t.Errorf("expected refreshed access token %q, got %q", "test_access_token_refreshed", token.AccessToken) + } +} + +func TestAuthorize_ScopeAccumulation(t *testing.T) { + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test_client_id": { + Secret: "test_client_secret", + RedirectURIs: []string{"http://localhost:12345/callback"}, + }, + }, + }, + TokenScopeFunc: func(requestedScope string) string { + // Simulate a server that never grants "write". + var granted []string + for _, s := range strings.Fields(requestedScope) { + if s != "write" { + granted = append(granted, s) + } + } + return strings.Join(granted, " ") + }, + }) + authServer.Start(t) + + resourceMux := http.NewServeMux() + resourceServer := httptest.NewServer(resourceMux) + t.Cleanup(resourceServer.Close) + resourceURL := resourceServer.URL + "/resource" + + resourceMux.Handle("/.well-known/oauth-protected-resource/resource", ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: resourceURL, + AuthorizationServers: []string{authServer.URL()}, + })) + + var capturedAuthURLs []string + noRedirectClient := &http.Client{ + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } + + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost:12345/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "test_client_secret", + }, + }, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + capturedAuthURLs = append(capturedAuthURLs, args.URL) + resp, err := noRedirectClient.Get(args.URL) + if err != nil { + return nil, err + } + defer resp.Body.Close() + loc, err := resp.Location() + if err != nil { + return nil, err + } + return &AuthorizationResult{ + Code: loc.Query().Get("code"), + State: loc.Query().Get("state"), + Iss: loc.Query().Get("iss"), + }, nil + }, + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + + // First authorization: 401 with scope="read write". + // The token response will only grant "read" (TokenScopeFunc strips "write"). + req := httptest.NewRequest(http.MethodGet, resourceURL, nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + } + resp.Header.Set("WWW-Authenticate", + fmt.Sprintf(`Bearer scope="read write", resource_metadata="%s/.well-known/oauth-protected-resource/resource"`, resourceServer.URL)) + if err := handler.Authorize(context.Background(), req, resp); err != nil { + t.Fatalf("First Authorize failed: %v", err) + } + + // Verify first auth URL requested "read" and "write". + firstURL, err := url.Parse(capturedAuthURLs[0]) + if err != nil { + t.Fatalf("Failed to parse first auth URL: %v", err) + } + firstScopes := strings.Fields(firstURL.Query().Get("scope")) + if diff := cmp.Diff([]string{"read", "write"}, firstScopes, cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" { + t.Errorf("First auth scopes mismatch (-want +got):\n%s", diff) + } + + // Verify only "read" was granted (the token omitted "write"). + issuer := authServer.URL() + if diff := cmp.Diff([]string{"read"}, handler.grantedScopes[issuer], cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" { + t.Errorf("After first Authorize, grantedScopes mismatch (-want +got):\n%s", diff) + } + + // Second authorization: 403 insufficient_scope with scope="admin". + // Accumulated scopes should be "read" (previously granted) + "admin" (new). + req2 := httptest.NewRequest(http.MethodGet, resourceURL, nil) + resp2 := &http.Response{ + StatusCode: http.StatusForbidden, + Header: make(http.Header), + Body: http.NoBody, + } + resp2.Header.Set("WWW-Authenticate", + fmt.Sprintf(`Bearer error="insufficient_scope", scope="admin", resource_metadata="%s/.well-known/oauth-protected-resource/resource"`, resourceServer.URL)) + if err := handler.Authorize(context.Background(), req2, resp2); err != nil { + t.Fatalf("Second Authorize failed: %v", err) + } + + // Verify second auth URL accumulated "read" (granted) + "admin" (challenged), + // but NOT "write" (requested but never granted). + secondURL, err := url.Parse(capturedAuthURLs[1]) + if err != nil { + t.Fatalf("Failed to parse second auth URL: %v", err) + } + secondScopes := strings.Fields(secondURL.Query().Get("scope")) + if diff := cmp.Diff([]string{"admin", "read"}, secondScopes, cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" { + t.Errorf("Second auth scopes mismatch (-want +got):\n%s", diff) + } +} + +func TestAuthorize_ForbiddenUnhandledError(t *testing.T) { + req := httptest.NewRequest(http.MethodGet, "http://example.com/resource", nil) + resp := &http.Response{ + StatusCode: http.StatusForbidden, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + resp.Header.Set( + "WWW-Authenticate", + "Bearer error=invalid_token", + ) + handler, err := NewAuthorizationCodeHandler(validConfig()) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + err = handler.Authorize(t.Context(), req, resp) + if err != nil { + t.Fatalf("Authorize() failed: %v", err) + } +} + +func TestNewAuthorizationCodeHandler_Success(t *testing.T) { + simpleHandler := func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + return nil, nil + } + tests := []struct { + name string + config *AuthorizationCodeHandlerConfig + }{ + { + name: "ClientIDMetadataDocumentConfig", + config: &AuthorizationCodeHandlerConfig{ + ClientIDMetadataDocumentConfig: &ClientIDMetadataDocumentConfig{URL: "https://example.com/client"}, + RedirectURL: "https://example.com/callback", + AuthorizationCodeFetcher: simpleHandler, + }, + }, + { + name: "PreregisteredClientConfig", + config: &AuthorizationCodeHandlerConfig{ + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "test_client_secret", + }, + }, + RedirectURL: "https://example.com/callback", + AuthorizationCodeFetcher: simpleHandler, + }, + }, + { + name: "DynamicClientRegistrationConfig_NoRedirectURL", + config: &AuthorizationCodeHandlerConfig{ + DynamicClientRegistrationConfig: &DynamicClientRegistrationConfig{ + Metadata: &oauthex.ClientRegistrationMetadata{ + RedirectURIs: []string{ + "https://example.com/callback", + }, + }, + }, + AuthorizationCodeFetcher: simpleHandler, + }, + }, + { + name: "DynamicClientRegistrationConfig_WithRedirectURL", + config: &AuthorizationCodeHandlerConfig{ + DynamicClientRegistrationConfig: &DynamicClientRegistrationConfig{ + Metadata: &oauthex.ClientRegistrationMetadata{ + RedirectURIs: []string{ + "https://example.com/callback", + }, + }, + }, + RedirectURL: "https://example.com/callback", + AuthorizationCodeFetcher: simpleHandler, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if _, err := NewAuthorizationCodeHandler(tt.config); err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + }) + } +} + +func TestNewAuthorizationCodeHandler_Error(t *testing.T) { + // Ensure the base config is valid. + if _, err := NewAuthorizationCodeHandler(validConfig()); err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + + tests := []struct { + name string + config func() *AuthorizationCodeHandlerConfig + }{ + { + name: "NilConfig", + config: func() *AuthorizationCodeHandlerConfig { + return nil + }, + }, + { + name: "NoRegistrationConfig", + config: func() *AuthorizationCodeHandlerConfig { + cfg := validConfig() + cfg.ClientIDMetadataDocumentConfig = nil + cfg.PreregisteredClient = nil + cfg.DynamicClientRegistrationConfig = nil + return cfg + }, + }, + { + name: "MissingRedirectURL", + config: func() *AuthorizationCodeHandlerConfig { + cfg := validConfig() + cfg.RedirectURL = "" + return cfg + }, + }, + { + name: "MissingAuthorizationCodeFetcher", + config: func() *AuthorizationCodeHandlerConfig { + cfg := validConfig() + cfg.AuthorizationCodeFetcher = nil + return cfg + }, + }, + { + name: "InvalidMetadataURL", + config: func() *AuthorizationCodeHandlerConfig { + cfg := validConfig() + cfg.ClientIDMetadataDocumentConfig.URL = "https://example.com" + return cfg + }, + }, + { + name: "InvalidPreregistered_MissingSecretConfig", + config: func() *AuthorizationCodeHandlerConfig { + cfg := validConfig() + cfg.PreregisteredClient = &oauthex.ClientCredentials{} + return cfg + }, + }, + { + name: "InvalidPreregistered_EmptyID", + config: func() *AuthorizationCodeHandlerConfig { + cfg := validConfig() + cfg.PreregisteredClient = &oauthex.ClientCredentials{ + ClientID: "", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "secret", + }, + } + return cfg + }, + }, + { + name: "InvalidPreregistered_EmptySecret", + config: func() *AuthorizationCodeHandlerConfig { + cfg := validConfig() + cfg.PreregisteredClient = &oauthex.ClientCredentials{ + ClientID: "test_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "", + }, + } + return cfg + }, + }, + { + name: "InvalidDynamic_MissingMetadata", + config: func() *AuthorizationCodeHandlerConfig { + cfg := validConfig() + cfg.DynamicClientRegistrationConfig = &DynamicClientRegistrationConfig{} + return cfg + }, + }, + { + name: "InvalidDynamic_InconsistentRedirectURI", + config: func() *AuthorizationCodeHandlerConfig { + cfg := validConfig() + cfg.DynamicClientRegistrationConfig = &DynamicClientRegistrationConfig{ + Metadata: &oauthex.ClientRegistrationMetadata{ + RedirectURIs: []string{"https://example.com/callback1"}, + }, + } + cfg.RedirectURL = "https://example.com/callback2" + return cfg + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := NewAuthorizationCodeHandler(tt.config()) + if err == nil { + t.Errorf("NewAuthorizationCodeHandler() = nil, want error") + } + }) + } +} + +func TestGetProtectedResourceMetadata_Success(t *testing.T) { + handler, err := NewAuthorizationCodeHandler(validConfig()) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler() error = %v", err) + } + + pathForChallenge := "/protected-resource" + + tests := []struct { + name string + challengesProvided bool + // Path of the PRM endpoint. + prmPath string + // Path of the MCP server that is accessed. + mcpServerPath string + // Path for the Resource expected in the returned PRM. + resourcePath string + }{ + { + name: "FromChallenges", + challengesProvided: true, + prmPath: pathForChallenge, + mcpServerPath: "/resource", + resourcePath: "/resource", + }, + { + name: "FallbackToEndpoint", + challengesProvided: false, + prmPath: "/.well-known/oauth-protected-resource/resource", + mcpServerPath: "/resource", + resourcePath: "/resource", + }, + { + name: "FallbackToRoot", + challengesProvided: false, + prmPath: "/.well-known/oauth-protected-resource", + mcpServerPath: "/resource", + resourcePath: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mux := http.NewServeMux() + server := httptest.NewServer(mux) + t.Cleanup(server.Close) + metadata := &oauthex.ProtectedResourceMetadata{ + Resource: server.URL + tt.resourcePath, + AuthorizationServers: []string{"https://oauth.example.com"}, + ScopesSupported: []string{"read", "write"}, + } + mux.Handle(tt.prmPath, ProtectedResourceMetadataHandler(metadata)) + var challenges []oauthex.Challenge + if tt.challengesProvided { + challenges = []oauthex.Challenge{ + { + Scheme: "Bearer", + Params: map[string]string{ + "resource_metadata": server.URL + pathForChallenge, + }, + }, + } + } + + got, err := handler.getProtectedResourceMetadata(t.Context(), challenges, server.URL+tt.mcpServerPath) + if err != nil { + t.Fatalf("getProtectedResourceMetadata() error = %v", err) + } + if got == nil { + t.Fatal("getProtectedResourceMetadata() got nil, want metadata") + } + if diff := cmp.Diff(metadata, got); diff != "" { + t.Errorf("getProtectedResourceMetadata() metadata mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestGetProtectedResourceMetadata_Backcompat(t *testing.T) { + handler, err := NewAuthorizationCodeHandler(validConfig()) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler() error = %v", err) + } + var challenges []oauthex.Challenge + got, err := handler.getProtectedResourceMetadata(t.Context(), challenges, "http://localhost:1234/resource") + if err != nil { + t.Fatalf("getProtectedResourceMetadata() error = %v", err) + } + wantPRM := &oauthex.ProtectedResourceMetadata{ + Resource: "http://localhost:1234/resource", + AuthorizationServers: []string{"http://localhost:1234"}, + } + if diff := cmp.Diff(wantPRM, got); diff != "" { + t.Errorf("getProtectedResourceMetadata() metadata mismatch (-want +got):\n%s", diff) + } +} + +func TestGetProtectedResourceMetadata_Error(t *testing.T) { + mux := http.NewServeMux() + server := httptest.NewServer(mux) + t.Cleanup(server.Close) + metadata := &oauthex.ProtectedResourceMetadata{ + Resource: server.URL + "/resource", + AuthorizationServers: nil, // Empty list is invalid + ScopesSupported: []string{"read", "write"}, + } + mux.Handle("/.well-known/oauth-protected-resource/resource", ProtectedResourceMetadataHandler(metadata)) + handler, err := NewAuthorizationCodeHandler(validConfig()) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler() error = %v", err) + } + var challenges []oauthex.Challenge + got, err := handler.getProtectedResourceMetadata(t.Context(), challenges, server.URL+"/resource") + if err == nil || !strings.Contains(err.Error(), "authorization servers") { + t.Errorf("getProtectedResourceMetadata() = %v, want error containing \"authorization servers\"", err) + } + if got != nil { + t.Errorf("getProtectedResourceMetadata() = %+v, want nil", got) + } +} + +func TestSelectTokenAuthMethod(t *testing.T) { + tests := []struct { + name string + supported []string + want oauth2.AuthStyle + }{ + { + name: "PostPreferredOverBasic", + supported: []string{"client_secret_basic", "client_secret_post"}, + want: oauth2.AuthStyleInParams, + }, + { + name: "BasicChosenIfPostNotSupported", + supported: []string{"private_key_jwt", "client_secret_basic"}, + want: oauth2.AuthStyleInHeader, + }, + { + name: "NoneSupported", + supported: []string{"private_key_jwt"}, + want: oauth2.AuthStyleAutoDetect, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := selectTokenAuthMethod(tt.supported) + if got != tt.want { + t.Errorf("selectTokenAuthMethod() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestHandleRegistration(t *testing.T) { + tests := []struct { + name string + serverConfig *oauthtest.RegistrationConfig + handlerConfig *AuthorizationCodeHandlerConfig + asm *oauthex.AuthServerMeta + want *resolvedClientConfig + wantError bool + issuerMatch bool + issuerSuffix string + }{ + { + name: "ClientIDMetadataDocument", + serverConfig: &oauthtest.RegistrationConfig{ + ClientIDMetadataDocumentSupported: true, + }, + handlerConfig: &AuthorizationCodeHandlerConfig{ + ClientIDMetadataDocumentConfig: &ClientIDMetadataDocumentConfig{URL: "https://client.example.com/metadata.json"}, + }, + want: &resolvedClientConfig{ + registrationType: registrationTypeClientIDMetadataDocument, + clientID: "https://client.example.com/metadata.json", + }, + }, + { + name: "Preregistered", + serverConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "pre_client_id": { + Secret: "pre_client_secret", + }, + }, + }, + handlerConfig: &AuthorizationCodeHandlerConfig{ + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "pre_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "pre_client_secret", + }, + }, + }, + want: &resolvedClientConfig{ + registrationType: registrationTypePreregistered, + clientID: "pre_client_id", + clientSecret: "pre_client_secret", + authStyle: oauth2.AuthStyleInParams, + }, + }, + { + name: "Preregistered_IssuerMatch", + serverConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "pre_client_id": { + Secret: "pre_client_secret", + }, + }, + }, + handlerConfig: &AuthorizationCodeHandlerConfig{ + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "pre_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "pre_client_secret", + }, + Issuer: "", // set dynamically in the test + }, + }, + want: &resolvedClientConfig{ + registrationType: registrationTypePreregistered, + clientID: "pre_client_id", + clientSecret: "pre_client_secret", + authStyle: oauth2.AuthStyleInParams, + }, + issuerMatch: true, + }, + { + name: "Preregistered_IssuerMismatch", + serverConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "pre_client_id": { + Secret: "pre_client_secret", + }, + }, + }, + handlerConfig: &AuthorizationCodeHandlerConfig{ + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "pre_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "pre_client_secret", + }, + Issuer: "https://other-issuer.example.com", + }, + }, + wantError: true, + }, + { + name: "Preregistered_IssuerMatchTrailingSlash", + serverConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "pre_client_id": { + Secret: "pre_client_secret", + }, + }, + }, + handlerConfig: &AuthorizationCodeHandlerConfig{ + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "pre_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "pre_client_secret", + }, + Issuer: "", // set dynamically in the test (with trailing slash) + }, + }, + want: &resolvedClientConfig{ + registrationType: registrationTypePreregistered, + clientID: "pre_client_id", + clientSecret: "pre_client_secret", + authStyle: oauth2.AuthStyleInParams, + }, + issuerMatch: true, + issuerSuffix: "/", + }, + { + name: "NoneSupported", + handlerConfig: &AuthorizationCodeHandlerConfig{ + ClientIDMetadataDocumentConfig: &ClientIDMetadataDocumentConfig{URL: "https://client.example.com/metadata.json"}, + }, + wantError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{RegistrationConfig: tt.serverConfig}) + s.Start(t) + // Set the Issuer dynamically if requested by the test case. + if tt.issuerMatch { + tt.handlerConfig.PreregisteredClient.Issuer = s.URL() + tt.issuerSuffix + } + tt.handlerConfig.AuthorizationCodeFetcher = func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + return nil, nil + } + tt.handlerConfig.RedirectURL = "https://example.com/callback" + handler, err := NewAuthorizationCodeHandler(tt.handlerConfig) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler() error = %v, want nil", err) + } + asm, err := GetAuthServerMetadata(t.Context(), s.URL(), http.DefaultClient) + if err != nil { + t.Fatalf("GetAuthServerMetadata() unexpected error = %v", err) + } + got, err := handler.handleRegistration(t.Context(), asm) + if err != nil { + if !tt.wantError { + t.Fatalf("handleRegistration() unexpected error = %v", err) + } + return + } + if tt.wantError { + t.Fatal("handleRegistration() expected error, got nil") + } + if got.registrationType != tt.want.registrationType { + t.Errorf("handleRegistration() registrationType = %v, want %v", got.registrationType, tt.want.registrationType) + } + if got.clientID != tt.want.clientID { + t.Errorf("handleRegistration() clientID = %q, want %q", got.clientID, tt.want.clientID) + } + if got.clientSecret != tt.want.clientSecret { + t.Errorf("handleRegistration() clientSecret = %q, want %q", got.clientSecret, tt.want.clientSecret) + } + if got.authStyle != tt.want.authStyle { + t.Errorf("handleRegistration() authStyle = %v, want %v", got.authStyle, tt.want.authStyle) + } + }) + } +} + +func TestDynamicRegistration(t *testing.T) { + s := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + RegistrationConfig: &oauthtest.RegistrationConfig{ + DynamicClientRegistrationEnabled: true, + }, + }) + s.Start(t) + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + DynamicClientRegistrationConfig: &DynamicClientRegistrationConfig{ + Metadata: &oauthex.ClientRegistrationMetadata{ + RedirectURIs: []string{"https://example.com/callback"}, + }, + }, + RedirectURL: "https://example.com/callback", + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + return nil, nil + }, + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler() error = %v", err) + } + asm, err := GetAuthServerMetadata(t.Context(), s.URL(), http.DefaultClient) + if err != nil { + t.Fatalf("GetAuthServerMetadata() unexpected error = %v", err) + } + got, err := handler.handleRegistration(t.Context(), asm) + if err != nil { + t.Fatalf("handleRegistration() error = %v, want nil", err) + } + if got.registrationType != registrationTypeDynamic { + t.Errorf("handleRegistration() registrationType = %v, want %v", got.registrationType, registrationTypeDynamic) + } + if got.clientID == "" { + t.Errorf("handleRegistration() clientID = %q, want non-empty", got.clientID) + } + if got.clientSecret == "" { + t.Errorf("handleRegistration() clientSecret = %q, want non-empty", got.clientSecret) + } + if got.authStyle != oauth2.AuthStyleInHeader { + t.Errorf("handleRegistration() authStyle = %v, want %v", got.authStyle, oauth2.AuthStyleInHeader) + } +} + +func TestValidateIssuerResponse(t *testing.T) { + const expectedIssuer = "https://auth.example.com" + + tests := []struct { + name string + iss string + issSupported bool + acceptUnadvertisedIss bool + wantErr bool + wantErrContains string + }{ + { + name: "ValidIss", + iss: expectedIssuer, + issSupported: true, + }, + { + name: "WrongIss", + iss: "https://attacker.example.com", + issSupported: true, + wantErr: true, + wantErrContains: "does not match expected issuer", + }, + { + name: "MissingIssWhenRequired", + iss: "", + issSupported: true, + wantErr: true, + wantErrContains: "RFC 9207", + }, + { + name: "MissingIssWhenNotRequired", + iss: "", + issSupported: false, + }, + { + // Default / zero-value policy: reject a matching unadvertised iss + // (released v1.7.0 behavior). + name: "UnadvertisedIssCorrectRejectedByDefault", + iss: expectedIssuer, + issSupported: false, + wantErr: true, + wantErrContains: "does not advertise", + }, + { + name: "UnadvertisedIssCorrectAcceptedWhenOptedIn", + iss: expectedIssuer, + issSupported: false, + acceptUnadvertisedIss: true, + }, + { + name: "UnadvertisedIssWrong", + iss: "https://attacker.example.com", + issSupported: false, + wantErr: true, + wantErrContains: "does not match expected issuer", + }, + { + name: "UnadvertisedIssWrongStillRejectedWhenOptedIn", + iss: "https://attacker.example.com", + issSupported: false, + acceptUnadvertisedIss: true, + wantErr: true, + wantErrContains: "does not match expected issuer", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateIssuerResponse(tt.iss, expectedIssuer, tt.issSupported, tt.acceptUnadvertisedIss) + if tt.wantErr { + if err == nil { + t.Fatalf("validateIssuerResponse() = nil, want error containing %q", tt.wantErrContains) + } + if !strings.Contains(err.Error(), tt.wantErrContains) { + t.Errorf("validateIssuerResponse() error = %q, want it to contain %q", err.Error(), tt.wantErrContains) + } + } else if err != nil { + t.Fatalf("validateIssuerResponse() unexpected error = %v", err) + } + }) + } +} + +// TestAuthorize_AcceptUnadvertisedIssPlumbing checks that Authorize passes +// AuthorizationCodeHandlerConfig.AcceptUnadvertisedIss through to +// validateIssuerResponse. The fake server returns iss without advertising +// authorization_response_iss_parameter_supported. +func TestAuthorize_AcceptUnadvertisedIssPlumbing(t *testing.T) { + tests := []struct { + name string + acceptUnadvertisedIss bool + wantErrContains string + }{ + { + name: "ZeroValueRejects", + wantErrContains: "does not advertise", + }, + { + name: "OptInAccepts", + acceptUnadvertisedIss: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + UnadvertiseIssParameter: true, + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test_client_id": { + Secret: "test_client_secret", + RedirectURIs: []string{"http://localhost:12345/callback"}, + }, + }, + }, + }) + authServer.Start(t) + + resourceMux := http.NewServeMux() + resourceServer := httptest.NewServer(resourceMux) + t.Cleanup(resourceServer.Close) + resourceURL := resourceServer.URL + "/resource" + resourceMux.Handle("/.well-known/oauth-protected-resource/resource", ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: resourceURL, + AuthorizationServers: []string{authServer.URL()}, + })) + + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost:12345/callback", + AcceptUnadvertisedIss: tt.acceptUnadvertisedIss, + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "test_client_secret", + }, + }, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + client := &http.Client{ + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } + resp, err := client.Get(args.URL) + if err != nil { + return nil, fmt.Errorf("failed to visit auth URL: %v", err) + } + defer resp.Body.Close() + location, err := resp.Location() + if err != nil { + return nil, fmt.Errorf("failed to get location header: %v", err) + } + return &AuthorizationResult{ + Code: location.Query().Get("code"), + State: location.Query().Get("state"), + Iss: location.Query().Get("iss"), + }, nil + }, + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, resourceURL, nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + resp.Header.Set( + "WWW-Authenticate", + "Bearer resource_metadata="+resourceServer.URL+"/.well-known/oauth-protected-resource/resource", + ) + err = handler.Authorize(context.Background(), req, resp) + if tt.wantErrContains != "" { + if err == nil { + t.Fatalf("Authorize() = nil, want error containing %q", tt.wantErrContains) + } + if !strings.Contains(err.Error(), tt.wantErrContains) { + t.Fatalf("Authorize() error = %q, want it to contain %q", err.Error(), tt.wantErrContains) + } + return + } + if err != nil { + t.Fatalf("Authorize() unexpected error = %v", err) + } + }) + } +} + +func TestInferApplicationType(t *testing.T) { + tests := []struct { + name string + redirectURIs []string + want string + }{ + { + name: "localhost", + redirectURIs: []string{"http://localhost:8085/callback"}, + want: "native", + }, + { + name: "127.0.0.1", + redirectURIs: []string{"http://127.0.0.1:8085/callback"}, + want: "native", + }, + { + name: "IPv6 loopback", + redirectURIs: []string{"http://[::1]:8085/callback"}, + want: "native", + }, + { + name: "custom scheme", + redirectURIs: []string{"myapp://callback"}, + want: "native", + }, + { + name: "HTTPS remote", + redirectURIs: []string{"https://myapp.example.com/callback"}, + want: "web", + }, + { + name: "mixed native and web", + redirectURIs: []string{"https://myapp.example.com/callback", "http://localhost:8085/callback"}, + want: "", + }, + { + name: "multiple remote", + redirectURIs: []string{"https://app1.example.com/cb", "https://app2.example.com/cb"}, + want: "web", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := inferApplicationType(tt.redirectURIs) + if got != tt.want { + t.Errorf("inferApplicationType() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestApplicationTypeInference(t *testing.T) { + fetcher := func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + return nil, nil + } + + tests := []struct { + name string + redirectURIs []string + initialAppType string + wantAppType string + wantErr bool + }{ + { + name: "inferred as native for localhost", + redirectURIs: []string{"http://localhost:8085/callback"}, + wantAppType: "native", + }, + { + name: "inferred as web for remote", + redirectURIs: []string{"https://example.com/callback"}, + wantAppType: "web", + }, + { + name: "mixed native and web URIs sets empty application type", + redirectURIs: []string{"https://example.com/callback", "http://localhost:8085/callback"}, + wantAppType: "", + }, + { + name: "explicit value matching inference is preserved", + redirectURIs: []string{"http://localhost:8085/callback"}, + initialAppType: "native", + wantAppType: "native", + }, + { + name: "explicit value conflicts with inference returns error", + redirectURIs: []string{"http://localhost:8085/callback"}, + initialAppType: "web", + wantErr: true, + }, + { + name: "explicit value when inference is ambiguous returns error", + redirectURIs: []string{"https://example.com/callback", "http://localhost:8085/callback"}, + initialAppType: "web", + wantErr: true, + }, + { + name: "invalid URI returns empty application type", + redirectURIs: []string{"http://%/"}, + wantAppType: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := &AuthorizationCodeHandlerConfig{ + DynamicClientRegistrationConfig: &DynamicClientRegistrationConfig{ + Metadata: &oauthex.ClientRegistrationMetadata{ + RedirectURIs: tt.redirectURIs, + ApplicationType: tt.initialAppType, + }, + }, + AuthorizationCodeFetcher: fetcher, + } + _, err := NewAuthorizationCodeHandler(cfg) + if (err != nil) != tt.wantErr { + t.Fatalf("NewAuthorizationCodeHandler() error = %v, wantErr %v", err, tt.wantErr) + } + if tt.wantErr { + return + } + got := cfg.DynamicClientRegistrationConfig.Metadata.ApplicationType + if got != tt.wantAppType { + t.Errorf("ApplicationType = %q, want %q", got, tt.wantAppType) + } + }) + } +} + +func TestAuthorize_OfflineAccessScope(t *testing.T) { + tests := []struct { + name string + requestRefreshToken bool + asScopesSupported []string + challengeScopes string + wantOfflineAccess bool + }{ + { + name: "AddedWhenASSupportsAndClientRequests", + requestRefreshToken: true, + asScopesSupported: []string{"openid", "offline_access"}, + wantOfflineAccess: true, + }, + { + name: "NotAddedWhenClientDoesNotRequest", + requestRefreshToken: false, + asScopesSupported: []string{"openid", "offline_access"}, + wantOfflineAccess: false, + }, + { + name: "NotAddedWhenASDoesNotSupport", + requestRefreshToken: true, + asScopesSupported: []string{"openid"}, + wantOfflineAccess: false, + }, + { + name: "NotDuplicatedWhenAlreadyInScopes", + requestRefreshToken: true, + asScopesSupported: []string{"openid", "offline_access"}, + challengeScopes: "read offline_access", + wantOfflineAccess: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + ScopesSupported: tt.asScopesSupported, + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test_client_id": { + Secret: "test_client_secret", + RedirectURIs: []string{"http://localhost:12345/callback"}, + }, + }, + }, + }) + authServer.Start(t) + + resourceMux := http.NewServeMux() + resourceServer := httptest.NewServer(resourceMux) + t.Cleanup(resourceServer.Close) + resourceURL := resourceServer.URL + "/resource" + resourceMux.Handle("/.well-known/oauth-protected-resource/resource", ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: resourceURL, + AuthorizationServers: []string{authServer.URL()}, + })) + + var capturedAuthURL string + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost:12345/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "test_client_secret", + }, + }, + RequestRefreshToken: tt.requestRefreshToken, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + capturedAuthURL = args.URL + return nil, fmt.Errorf("stop after capturing URL") + }, + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, resourceURL, nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + wwwAuth := "Bearer resource_metadata=" + resourceServer.URL + "/.well-known/oauth-protected-resource/resource" + if tt.challengeScopes != "" { + wwwAuth += fmt.Sprintf(", scope=%q", tt.challengeScopes) + } + resp.Header.Set("WWW-Authenticate", wwwAuth) + + // Authorize will fail at the fetcher, but we only care about the URL. + handler.Authorize(context.Background(), req, resp) + + if capturedAuthURL == "" { + t.Fatal("AuthorizationCodeFetcher was not called") + } + u, err := url.Parse(capturedAuthURL) + if err != nil { + t.Fatalf("failed to parse captured auth URL: %v", err) + } + scopes := strings.Fields(u.Query().Get("scope")) + hasOfflineAccess := slices.Contains(scopes, "offline_access") + if hasOfflineAccess != tt.wantOfflineAccess { + t.Errorf("offline_access in scopes = %v, want %v (scopes: %v)", hasOfflineAccess, tt.wantOfflineAccess, scopes) + } + + // When offline_access was already present in challenge scopes, + // verify it appears exactly once. + if tt.wantOfflineAccess { + count := 0 + for _, s := range scopes { + if s == "offline_access" { + count++ + } + } + if count != 1 { + t.Errorf("offline_access appears %d times in scopes, want 1", count) + } + } + }) + } +} + +func TestAuthorize_ScopeFilter(t *testing.T) { + // advertised is delivered via the WWW-Authenticate "scope" challenge, which + // the handler treats as the discovered scopes passed to ScopeFilter. + const advertised = "gmail.metadata gmail.readonly gmail.compose" + tests := []struct { + name string + filter func(discovered []string) []string + want []string // requested scopes, order-independent + }{ + { + name: "nil filter leaves scopes unchanged", + filter: nil, + want: []string{"gmail.metadata", "gmail.readonly", "gmail.compose"}, + }, + { + name: "filter drops a scope", + filter: func(d []string) []string { + return slices.DeleteFunc(slices.Clone(d), func(s string) bool { return s == "gmail.metadata" }) + }, + want: []string{"gmail.readonly", "gmail.compose"}, + }, + { + name: "filter replaces the set entirely", + filter: func([]string) []string { return []string{"gmail.readonly"} }, + want: []string{"gmail.readonly"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test_client_id": { + Secret: "test_client_secret", + RedirectURIs: []string{"http://localhost:12345/callback"}, + }, + }, + }, + }) + authServer.Start(t) + + resourceMux := http.NewServeMux() + resourceServer := httptest.NewServer(resourceMux) + t.Cleanup(resourceServer.Close) + resourceURL := resourceServer.URL + "/resource" + resourceMux.Handle("/.well-known/oauth-protected-resource/resource", ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: resourceURL, + AuthorizationServers: []string{authServer.URL()}, + })) + + var capturedAuthURL string + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost:12345/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ClientSecret: "test_client_secret"}, + }, + ScopeFilter: tt.filter, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + capturedAuthURL = args.URL + return nil, fmt.Errorf("stop after capturing URL") + }, + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, resourceURL, nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + resp.Header.Set("WWW-Authenticate", fmt.Sprintf( + "Bearer resource_metadata=%s/.well-known/oauth-protected-resource/resource, scope=%q", + resourceServer.URL, advertised)) + + handler.Authorize(context.Background(), req, resp) + + if capturedAuthURL == "" { + t.Fatal("AuthorizationCodeFetcher was not called") + } + u, err := url.Parse(capturedAuthURL) + if err != nil { + t.Fatalf("failed to parse captured auth URL: %v", err) + } + // Compare as a set: UnionScopes (applied downstream) returns map keys, + // so the order of the requested scope parameter is not deterministic. + got := strings.Fields(u.Query().Get("scope")) + want := slices.Clone(tt.want) + slices.Sort(got) + slices.Sort(want) + if !slices.Equal(got, want) { + t.Errorf("requested scopes = %v, want %v (any order)", got, want) + } + }) + } +} + +// validConfig for test to create an AuthorizationCodeHandler using its constructor. +// Values that are relevant to the test should be set explicitly. +func validConfig() *AuthorizationCodeHandlerConfig { + return &AuthorizationCodeHandlerConfig{ + ClientIDMetadataDocumentConfig: &ClientIDMetadataDocumentConfig{URL: "https://example.com/client"}, + RedirectURL: "https://example.com/callback", + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + return nil, nil + }, + } +} + +func TestNewTokenSource(t *testing.T) { + // mock the /token endpoint to successfully return an access token on code exchange + mockTS := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/token" { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"access_token": "test_token", "token_type": "bearer"}`)) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer mockTS.Close() + + // configure the handler and set NewTokenSource + var called bool + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client", + }, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + u, _ := url.Parse(args.URL) + return &AuthorizationResult{ + Code: "test_code", + State: u.Query().Get("state"), + }, nil + }, + NewTokenSource: func(ctx context.Context, cfg *oauth2.Config, token *oauth2.Token) (oauth2.TokenSource, error) { + called = true + return oauth2.StaticTokenSource(token), nil + }, + Client: mockTS.Client(), + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Simulate a 401 response from a resource server. + // The WWW-Authenticate: Bearer header triggers the authorization logic. + req := httptest.NewRequest(http.MethodGet, mockTS.URL, nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + resp.Header.Set("WWW-Authenticate", "Bearer") + + // Authorize and confirm NewTokenSource was called. + err = handler.Authorize(t.Context(), req, resp) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !called { + t.Error("expected NewTokenSource to be called") + } +} + +func TestInitialTokenSource(t *testing.T) { + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost:12345/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client_id", + }, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + return nil, nil + }, + InitialTokenSource: oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "set_token"}), + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + + ts, err := handler.TokenSource(t.Context()) + if err != nil { + t.Fatalf("failed to get token source: %v", err) + } + if ts == nil { + t.Fatal("expected token source to be non-nil") + } + + tok, err := ts.Token() + if err != nil { + t.Fatalf("failed to get Token: %v", err) + } + if tok.AccessToken != "set_token" { + t.Errorf("expected access token 'set_token', got '%s'", tok.AccessToken) + } +} diff --git a/auth/client.go b/auth/client.go index 3ed6b69..db32d97 100644 --- a/auth/client.go +++ b/auth/client.go @@ -2,140 +2,39 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. -//go:build mcp_go_client_oauth - package auth import ( - "bytes" "context" - "errors" - "io" "net/http" - "sync" - "github.com/modelcontextprotocol/go-sdk/oauthex" "golang.org/x/oauth2" ) -// An OAuthHandler conducts an OAuth flow and returns a [oauth2.TokenSource] if the authorization -// is approved, or an error if not. -type OAuthHandler func(context.Context, OAuthHandlerArgs) (oauth2.TokenSource, error) +// OAuthHandler is an interface for handling OAuth flows. +// +// If a transport wishes to support OAuth 2 authorization, it should support +// being configured with an OAuthHandler. It should call the handler's +// TokenSource method whenever it sends an HTTP request to set the +// Authorization header. If a request fails with a 401 or 403, it should call +// Authorize, and if that returns nil, it should retry the request. It should +// not call Authorize after the second failure. See +// [github.com/modelcontextprotocol/go-sdk/mcp.StreamableClientTransport] +// for an example. +type OAuthHandler interface { + // TokenSource returns a token source to be used for outgoing requests. + // Returned token source might be nil. In that case, the transport will not + // add any authorization headers to the request. + TokenSource(context.Context) (oauth2.TokenSource, error) -// OAuthHandlerArgs are arguments to an [OAuthHandler]. -type OAuthHandlerArgs struct { - // The URL to fetch protected resource metadata, extracted from the WWW-Authenticate header. - // Empty if not present or there was an error obtaining it. - ResourceMetadataURL string -} - -// HTTPTransport is an [http.RoundTripper] that follows the MCP -// OAuth protocol when it encounters a 401 Unauthorized response. -type HTTPTransport struct { - handler OAuthHandler - mu sync.Mutex // protects opts.Base - opts HTTPTransportOptions -} - -// NewHTTPTransport returns a new [*HTTPTransport]. -// The handler is invoked when an HTTP request results in a 401 Unauthorized status. -// It is called only once per transport. Once a TokenSource is obtained, it is used -// for the lifetime of the transport; subsequent 401s are not processed. -func NewHTTPTransport(handler OAuthHandler, opts *HTTPTransportOptions) (*HTTPTransport, error) { - if handler == nil { - return nil, errors.New("handler cannot be nil") - } - t := &HTTPTransport{ - handler: handler, - } - if opts != nil { - t.opts = *opts - } - if t.opts.Base == nil { - t.opts.Base = http.DefaultTransport - } - return t, nil -} - -// HTTPTransportOptions are options to [NewHTTPTransport]. -type HTTPTransportOptions struct { - // Base is the [http.RoundTripper] to use. - // If nil, [http.DefaultTransport] is used. - Base http.RoundTripper -} - -func (t *HTTPTransport) RoundTrip(req *http.Request) (*http.Response, error) { - t.mu.Lock() - base := t.opts.Base - t.mu.Unlock() - - var ( - // If haveBody is set, the request has a nontrivial body, and we need avoid - // reading (or closing) it multiple times. In that case, bodyBytes is its - // content. - haveBody bool - bodyBytes []byte - ) - if req.Body != nil && req.Body != http.NoBody { - // if we're setting Body, we must mutate first. - req = req.Clone(req.Context()) - haveBody = true - var err error - bodyBytes, err = io.ReadAll(req.Body) - if err != nil { - return nil, err - } - // Now that we've read the request body, http.RoundTripper requires that we - // close it. - req.Body.Close() // ignore error - req.Body = io.NopCloser(bytes.NewReader(bodyBytes)) - } - - resp, err := base.RoundTrip(req) - if err != nil { - return nil, err - } - if resp.StatusCode != http.StatusUnauthorized { - return resp, nil - } - if _, ok := base.(*oauth2.Transport); ok { - // We failed to authorize even with a token source; give up. - return resp, nil - } - - resp.Body.Close() - // Try to authorize. - t.mu.Lock() - defer t.mu.Unlock() - // If we don't have a token source, get one by following the OAuth flow. - // (We may have obtained one while t.mu was not held above.) - // TODO: We hold the lock for the entire OAuth flow. This could be a long - // time. Is there a better way? - if _, ok := t.opts.Base.(*oauth2.Transport); !ok { - authHeaders := resp.Header[http.CanonicalHeaderKey("WWW-Authenticate")] - ts, err := t.handler(req.Context(), OAuthHandlerArgs{ - ResourceMetadataURL: extractResourceMetadataURL(authHeaders), - }) - if err != nil { - return nil, err - } - t.opts.Base = &oauth2.Transport{Base: t.opts.Base, Source: ts} - } - - // If we don't have a body, the request is reusable, though it will be cloned - // by the base. However, if we've had to read the body, we must clone. - if haveBody { - req = req.Clone(req.Context()) - req.Body = io.NopCloser(bytes.NewReader(bodyBytes)) - } - - return t.opts.Base.RoundTrip(req) -} - -func extractResourceMetadataURL(authHeaders []string) string { - cs, err := oauthex.ParseWWWAuthenticate(authHeaders) - if err != nil { - return "" - } - return oauthex.ResourceMetadataURL(cs) + // Authorize is called when an HTTP request results in an error that may + // be addressed by the authorization flow (currently 401 Unauthorized and 403 Forbidden). + // It is responsible for performing the OAuth flow to obtain an access token. + // The arguments are the request that failed and the response that was received for it. + // The headers of the request are available, but the body will have already been consumed + // when Authorize is called. + // If the returned error is nil, TokenSource is expected to return a non-nil token source. + // After a successful call to Authorize, the HTTP request will be retried by the transport. + // The function is responsible for closing the response body. + Authorize(context.Context, *http.Request, *http.Response) error } diff --git a/auth/client_test.go b/auth/client_test.go deleted file mode 100644 index e1b7dc7..0000000 --- a/auth/client_test.go +++ /dev/null @@ -1,152 +0,0 @@ -// Copyright 2025 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -//go:build mcp_go_client_oauth - -package auth - -import ( - "context" - "errors" - "fmt" - "io" - "net/http" - "net/http/httptest" - "strings" - "testing" - - "golang.org/x/oauth2" -) - -// A basicReader is an io.Reader to be used as a non-rereadable request body. -// -// net/http has special handling for strings.Reader that we want to avoid. -type basicReader struct { - r *strings.Reader -} - -func (r *basicReader) Read(p []byte) (n int, err error) { return r.r.Read(p) } - -// TestHTTPTransport validates the OAuth HTTPTransport. -func TestHTTPTransport(t *testing.T) { - const testToken = "test-token-123" - fakeTokenSource := oauth2.StaticTokenSource(&oauth2.Token{ - AccessToken: testToken, - TokenType: "Bearer", - }) - - // authServer simulates a resource that requires OAuth. - authServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodPost { - // Ensure that the body was properly cloned, by reading it completely. - // If the body is not cloned, reading it the second time may yield no - // bytes. - body, err := io.ReadAll(r.Body) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - if len(body) == 0 { - http.Error(w, "empty body", http.StatusBadRequest) - return - } - } - authHeader := r.Header.Get("Authorization") - if authHeader == fmt.Sprintf("Bearer %s", testToken) { - w.WriteHeader(http.StatusOK) - return - } - - w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="http://metadata.example.com"`) - w.WriteHeader(http.StatusUnauthorized) - })) - defer authServer.Close() - - t.Run("successful auth flow", func(t *testing.T) { - var handlerCalls int - handler := func(ctx context.Context, args OAuthHandlerArgs) (oauth2.TokenSource, error) { - handlerCalls++ - if args.ResourceMetadataURL != "http://metadata.example.com" { - t.Errorf("handler got metadata URL %q, want %q", args.ResourceMetadataURL, "http://metadata.example.com") - } - return fakeTokenSource, nil - } - - transport, err := NewHTTPTransport(handler, nil) - if err != nil { - t.Fatalf("NewHTTPTransport() failed: %v", err) - } - client := &http.Client{Transport: transport} - - resp, err := client.Get(authServer.URL) - if err != nil { - t.Fatalf("client.Get() failed: %v", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - t.Errorf("got status %d, want %d", resp.StatusCode, http.StatusOK) - } - if handlerCalls != 1 { - t.Errorf("handler was called %d times, want 1", handlerCalls) - } - - // Second request should reuse the token and not call the handler again. - resp2, err := client.Get(authServer.URL) - if err != nil { - t.Fatalf("second client.Get() failed: %v", err) - } - defer resp2.Body.Close() - - if resp2.StatusCode != http.StatusOK { - t.Errorf("second request got status %d, want %d", resp2.StatusCode, http.StatusOK) - } - if handlerCalls != 1 { - t.Errorf("handler should still be called only once, but was %d", handlerCalls) - } - }) - - t.Run("request body is cloned", func(t *testing.T) { - handler := func(ctx context.Context, args OAuthHandlerArgs) (oauth2.TokenSource, error) { - if args.ResourceMetadataURL != "http://metadata.example.com" { - t.Errorf("handler got metadata URL %q, want %q", args.ResourceMetadataURL, "http://metadata.example.com") - } - return fakeTokenSource, nil - } - - transport, err := NewHTTPTransport(handler, nil) - if err != nil { - t.Fatalf("NewHTTPTransport() failed: %v", err) - } - client := &http.Client{Transport: transport} - - resp, err := client.Post(authServer.URL, "application/json", &basicReader{strings.NewReader("{}")}) - if err != nil { - t.Fatalf("client.Post() failed: %v", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - t.Errorf("got status %d, want %d", resp.StatusCode, http.StatusOK) - } - }) - - t.Run("handler returns error", func(t *testing.T) { - handlerErr := errors.New("user rejected auth") - handler := func(ctx context.Context, args OAuthHandlerArgs) (oauth2.TokenSource, error) { - return nil, handlerErr - } - - transport, err := NewHTTPTransport(handler, nil) - if err != nil { - t.Fatalf("NewHTTPTransport() failed: %v", err) - } - client := &http.Client{Transport: transport} - - _, err = client.Get(authServer.URL) - if !errors.Is(err, handlerErr) { - t.Errorf("client.Get() returned error %v, want %v", err, handlerErr) - } - }) -} diff --git a/auth/extauth/client_credentials.go b/auth/extauth/client_credentials.go new file mode 100644 index 0000000..c65d8b0 --- /dev/null +++ b/auth/extauth/client_credentials.go @@ -0,0 +1,279 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package extauth + +import ( + "context" + "fmt" + "io" + "net/http" + "net/url" + "slices" + "strings" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/internal/authutil" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" + "golang.org/x/oauth2/clientcredentials" +) + +// ClientCredentialsHandlerConfig is the configuration for [ClientCredentialsHandler]. +type ClientCredentialsHandlerConfig struct { + // Credentials contains the pre-registered client ID and secret. + // REQUIRED. Both ClientID and ClientSecretAuth must be set, since the + // client credentials grant requires a confidential client. + Credentials *oauthex.ClientCredentials + + // HTTPClient is an optional HTTP client for customization. + // If nil, http.DefaultClient is used. + // OPTIONAL. + HTTPClient *http.Client +} + +// ClientCredentialsHandler is an implementation of [auth.OAuthHandler] that +// uses the OAuth 2.0 Client Credentials grant (RFC 6749 Section 4.4) to +// obtain access tokens. +// +// This handler is intended for service-to-service authentication where the +// client has pre-registered credentials (client ID and secret) and does not +// require user interaction. It bypasses both dynamic client registration and +// the authorization code flow. +// +// The token endpoint and scopes are discovered automatically via Protected +// Resource Metadata (RFC 9728) and Authorization Server Metadata (RFC 8414), +// following the ext-auth specification SEP-1046. +type ClientCredentialsHandler struct { + config *ClientCredentialsHandlerConfig + tokenSource oauth2.TokenSource + + // grantedScopes maps authorization server issuer to the list of scopes granted by that issuer. + grantedScopes map[string][]string +} + +// Compile-time check that ClientCredentialsHandler implements auth.OAuthHandler. +var _ auth.OAuthHandler = (*ClientCredentialsHandler)(nil) + +// NewClientCredentialsHandler creates a new ClientCredentialsHandler. +// It validates the configuration and returns an error if invalid. +func NewClientCredentialsHandler(config *ClientCredentialsHandlerConfig) (*ClientCredentialsHandler, error) { + if config == nil { + return nil, fmt.Errorf("config must be provided") + } + if config.Credentials == nil { + return nil, fmt.Errorf("credentials are required") + } + if err := config.Credentials.Validate(); err != nil { + return nil, fmt.Errorf("invalid credentials: %w", err) + } + if config.Credentials.ClientSecretAuth == nil { + return nil, fmt.Errorf("clientSecretAuth is required for client credentials grant") + } + return &ClientCredentialsHandler{ + config: config, + grantedScopes: make(map[string][]string), + }, nil +} + +// TokenSource returns the token source for outgoing requests. +// Returns nil if authorization has not been performed yet. +func (h *ClientCredentialsHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) { + return h.tokenSource, nil +} + +// Authorize performs the Client Credentials grant to obtain an access token. +// It is called when a request fails with 401 or 403. +// +// The flow follows the ext-auth specification SEP-1046: +// 1. Discover Protected Resource Metadata from the request URL +// 2. Discover Authorization Server Metadata from PRM +// 3. Exchange client credentials for an access token at the token endpoint +func (h *ClientCredentialsHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error { + defer resp.Body.Close() + defer io.Copy(io.Discard, resp.Body) + + httpClient := h.config.HTTPClient + if httpClient == nil { + httpClient = http.DefaultClient + } + + // Step 1: Discover Protected Resource Metadata. + wwwChallenges, err := oauthex.ParseWWWAuthenticate(resp.Header[http.CanonicalHeaderKey("WWW-Authenticate")]) + if err != nil { + return fmt.Errorf("failed to parse WWW-Authenticate header: %v", err) + } + + prm, err := getProtectedResourceMetadata(ctx, wwwChallenges, req.URL.String(), httpClient) + if err != nil { + return err + } + + if len(prm.AuthorizationServers) == 0 { + return fmt.Errorf("protected resource metadata has no authorization servers specified") + } + + // Step 2: Discover Authorization Server Metadata. + asm, err := auth.GetAuthServerMetadata(ctx, prm.AuthorizationServers[0], httpClient) + if err != nil { + return fmt.Errorf("failed to get authorization server metadata: %w", err) + } + if asm == nil { + // Fallback to 2025-03-26 spec: predefined endpoints. + authServerURL := prm.AuthorizationServers[0] + asm = &oauthex.AuthServerMeta{ + Issuer: authServerURL, + TokenEndpoint: authServerURL + "/token", + } + } + + creds := h.config.Credentials + if creds.Issuer != "" && !authutil.IssuersEqual(creds.Issuer, asm.Issuer) { + return fmt.Errorf("authorization server issuer %q does not match pre-registered credentials issuer %q", asm.Issuer, creds.Issuer) + } + + // Determine requestedScopes: use PRM's scopes_supported if available. + requestedScopes := scopesFromChallenges(wwwChallenges) + if len(requestedScopes) == 0 && len(prm.ScopesSupported) > 0 { + requestedScopes = prm.ScopesSupported + } + + // Accumulate scopes: union previously granted scopes with the newly + // challenged scopes so that step-up authorization does not lose + // permissions granted in earlier rounds (SEP-2350). + requestedScopes = authutil.UnionScopes(h.grantedScopes[asm.Issuer], requestedScopes) + + // Step 3: Exchange client credentials for an access token. + cfg := &clientcredentials.Config{ + ClientID: creds.ClientID, + ClientSecret: creds.ClientSecretAuth.ClientSecret, + TokenURL: asm.TokenEndpoint, + Scopes: requestedScopes, + AuthStyle: selectTokenAuthMethod(asm.TokenEndpointAuthMethodsSupported), + } + + ctxWithClient := context.WithValue(ctx, oauth2.HTTPClient, httpClient) + h.tokenSource = cfg.TokenSource(ctxWithClient) + + return h.updateGrantedScopes(asm.Issuer, requestedScopes) +} + +func (h *ClientCredentialsHandler) updateGrantedScopes(issuer string, requestedScopes []string) error { + if h.tokenSource == nil { + return nil + } + tok, err := h.tokenSource.Token() + if err != nil { + h.tokenSource = nil + return fmt.Errorf("client credentials token request failed: %w", err) + } + if tokenScopes := authutil.ScopesFromToken(tok); tokenScopes == nil { + h.grantedScopes[issuer] = requestedScopes + } else { + h.grantedScopes[issuer] = tokenScopes + } + return nil +} + +// getProtectedResourceMetadata discovers Protected Resource Metadata (RFC 9728) +// from the request URL. This mirrors the logic in AuthorizationCodeHandler. +func getProtectedResourceMetadata(ctx context.Context, wwwChallenges []oauthex.Challenge, mcpServerURL string, httpClient *http.Client) (*oauthex.ProtectedResourceMetadata, error) { + // Use MCP server URL as the resource URI per + // https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#canonical-server-uri. + for _, u := range protectedResourceMetadataURLs(resourceMetadataURLFromChallenges(wwwChallenges), mcpServerURL) { + prm, err := oauthex.GetProtectedResourceMetadata(ctx, u.url, u.resource, httpClient) + if err != nil { + continue + } + if prm == nil { + continue + } + if len(prm.AuthorizationServers) == 0 { + // If we found PRM, we enforce the 2025-11-25 spec and not search further. + return nil, fmt.Errorf("protected resource metadata has no authorization servers specified") + } + return prm, nil + } + // Fallback to 2025-03-26 spec: MCP server root is the Authorization Server. + // https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization#server-metadata-discovery + u, err := url.Parse(mcpServerURL) + if err != nil { + return nil, fmt.Errorf("failed to parse MCP server URL: %v", err) + } + u.Path = "" + return &oauthex.ProtectedResourceMetadata{ + AuthorizationServers: []string{u.String()}, + Resource: mcpServerURL, + }, nil +} + +type prmURL struct { + url string + resource string +} + +// protectedResourceMetadataURLs returns URLs to try for PRM discovery. +// This mirrors the logic in AuthorizationCodeHandler. +func protectedResourceMetadataURLs(metadataURL, resourceURL string) []prmURL { + var urls []prmURL + if metadataURL != "" { + urls = append(urls, prmURL{url: metadataURL, resource: resourceURL}) + } + ru, err := url.Parse(resourceURL) + if err != nil { + return urls + } + mu := *ru + // At the path of the server's MCP endpoint. + mu.Path = "/.well-known/oauth-protected-resource/" + strings.TrimLeft(ru.Path, "/") + urls = append(urls, prmURL{url: mu.String(), resource: resourceURL}) + // At the root. + mu.Path = "/.well-known/oauth-protected-resource" + ru.Path = "" + urls = append(urls, prmURL{url: mu.String(), resource: ru.String()}) + return urls +} + +// resourceMetadataURLFromChallenges returns a resource metadata URL from +// WWW-Authenticate challenges, or the empty string if there is none. +func resourceMetadataURLFromChallenges(cs []oauthex.Challenge) string { + for _, c := range cs { + if u := c.Params["resource_metadata"]; u != "" { + return u + } + } + return "" +} + +// scopesFromChallenges returns scopes from WWW-Authenticate challenges. +// It only looks at challenges with the "Bearer" scheme. +func scopesFromChallenges(cs []oauthex.Challenge) []string { + for _, c := range cs { + if c.Scheme == "bearer" && c.Params["scope"] != "" { + return strings.Fields(c.Params["scope"]) + } + } + return nil +} + +// selectTokenAuthMethod selects the preferred token endpoint auth method based on +// the authorization server's supported methods. Prefers client_secret_post over +// client_secret_basic per the OAuth 2.1 draft. +func selectTokenAuthMethod(supported []string) oauth2.AuthStyle { + prefOrder := []string{ + "client_secret_post", + "client_secret_basic", + } + for _, method := range prefOrder { + if slices.Contains(supported, method) { + switch method { + case "client_secret_post": + return oauth2.AuthStyleInParams + case "client_secret_basic": + return oauth2.AuthStyleInHeader + } + } + } + return oauth2.AuthStyleAutoDetect +} diff --git a/auth/extauth/client_credentials_test.go b/auth/extauth/client_credentials_test.go new file mode 100644 index 0000000..b2973de --- /dev/null +++ b/auth/extauth/client_credentials_test.go @@ -0,0 +1,420 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package extauth + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/internal/oauthtest" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" +) + +func validClientCredentialsConfig() *ClientCredentialsHandlerConfig { + return &ClientCredentialsHandlerConfig{ + Credentials: &oauthex.ClientCredentials{ + ClientID: "test-client", + ClientSecretAuth: &oauthex.ClientSecretAuth{ClientSecret: "test-secret"}, + }, + } +} + +func TestNewClientCredentialsHandler_Validation(t *testing.T) { + tests := []struct { + name string + config *ClientCredentialsHandlerConfig + wantError string + }{ + { + name: "nil config", + config: nil, + wantError: "config must be provided", + }, + { + name: "nil credentials", + config: func() *ClientCredentialsHandlerConfig { + c := validClientCredentialsConfig() + c.Credentials = nil + return c + }(), + wantError: "credentials are required", + }, + { + name: "missing client ID", + config: func() *ClientCredentialsHandlerConfig { + c := validClientCredentialsConfig() + c.Credentials.ClientID = "" + return c + }(), + wantError: "ClientID is required", + }, + { + name: "missing client secret auth", + config: func() *ClientCredentialsHandlerConfig { + c := validClientCredentialsConfig() + c.Credentials.ClientSecretAuth = nil + return c + }(), + wantError: "clientSecretAuth is required", + }, + { + name: "empty client secret", + config: func() *ClientCredentialsHandlerConfig { + c := validClientCredentialsConfig() + c.Credentials.ClientSecretAuth.ClientSecret = "" + return c + }(), + wantError: "ClientSecret is required", + }, + { + name: "valid config", + config: validClientCredentialsConfig(), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, err := NewClientCredentialsHandler(tc.config) + if tc.wantError != "" { + if err == nil { + t.Errorf("expected error containing %q, got nil", tc.wantError) + } else if !strings.Contains(err.Error(), tc.wantError) { + t.Errorf("error %q does not contain %q", err.Error(), tc.wantError) + } + } else if err != nil { + t.Errorf("unexpected error: %v", err) + } + }) + } +} + +func TestClientCredentialsHandler_Authorize(t *testing.T) { + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + MetadataEndpointConfig: &oauthtest.MetadataEndpointConfig{ + ServeOAuthInsertedEndpoint: true, + }, + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test-client": {Secret: "test-secret"}, + }, + }, + ClientCredentialsConfig: &oauthtest.ClientCredentialsConfig{ + Enabled: true, + }, + }) + authServer.Start(t) + + resourceMux := http.NewServeMux() + resourceServer := httptest.NewServer(resourceMux) + t.Cleanup(resourceServer.Close) + resourceURL := resourceServer.URL + "/resource" + + resourceMux.Handle("/.well-known/oauth-protected-resource/resource", auth.ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: resourceURL, + AuthorizationServers: []string{authServer.URL()}, + ScopesSupported: []string{"mcp:read", "mcp:write"}, + })) + + t.Run("successful authorization", func(t *testing.T) { + handler, err := NewClientCredentialsHandler(validClientCredentialsConfig()) + if err != nil { + t.Fatal(err) + } + + // TokenSource should be nil before Authorize. + ts, err := handler.TokenSource(t.Context()) + if err != nil { + t.Fatal(err) + } + if ts != nil { + t.Fatal("expected nil TokenSource before Authorize") + } + + // Simulate a 401 response from the MCP server. + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: http.Header{}, + Body: http.NoBody, + } + req := httptest.NewRequest("GET", resourceURL, nil) + if err := handler.Authorize(t.Context(), req, resp); err != nil { + t.Fatalf("Authorize failed: %v", err) + } + + // TokenSource should now return a valid token. + ts, err = handler.TokenSource(t.Context()) + if err != nil { + t.Fatal(err) + } + if ts == nil { + t.Fatal("expected non-nil TokenSource after Authorize") + } + tok, err := ts.Token() + if err != nil { + t.Fatalf("Token() failed: %v", err) + } + if tok.AccessToken != "test_access_token" { + t.Errorf("got access token %q, want %q", tok.AccessToken, "test_access_token") + } + }) + + t.Run("bad credentials", func(t *testing.T) { + config := validClientCredentialsConfig() + config.Credentials.ClientSecretAuth.ClientSecret = "wrong-secret" + handler, err := NewClientCredentialsHandler(config) + if err != nil { + t.Fatal(err) + } + + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: http.Header{}, + Body: http.NoBody, + } + req := httptest.NewRequest("GET", resourceURL, nil) + if err := handler.Authorize(t.Context(), req, resp); err == nil { + t.Error("expected Authorize to fail with bad credentials") + } + + // TokenSource should still be nil after failed Authorize. + ts, err := handler.TokenSource(t.Context()) + if err != nil { + t.Fatal(err) + } + if ts != nil { + t.Error("expected nil TokenSource after failed Authorize") + } + }) + + t.Run("scopes from WWW-Authenticate challenge", func(t *testing.T) { + handler, err := NewClientCredentialsHandler(validClientCredentialsConfig()) + if err != nil { + t.Fatal(err) + } + + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: http.Header{ + "WWW-Authenticate": []string{`Bearer scope="read write"`}, + }, + Body: http.NoBody, + } + req := httptest.NewRequest("GET", resourceURL, nil) + if err := handler.Authorize(t.Context(), req, resp); err != nil { + t.Fatalf("Authorize failed: %v", err) + } + + ts, err := handler.TokenSource(t.Context()) + if err != nil { + t.Fatal(err) + } + if ts == nil { + t.Fatal("expected non-nil TokenSource") + } + }) + + t.Run("issuer mismatch", func(t *testing.T) { + config := validClientCredentialsConfig() + config.Credentials.Issuer = "https://other-issuer.example.com" + handler, err := NewClientCredentialsHandler(config) + if err != nil { + t.Fatal(err) + } + + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: http.Header{}, + Body: http.NoBody, + } + req := httptest.NewRequest("GET", resourceURL, nil) + err = handler.Authorize(t.Context(), req, resp) + if err == nil { + t.Fatal("expected Authorize to fail with issuer mismatch") + } + if !strings.Contains(err.Error(), "does not match") { + t.Errorf("error %q does not mention issuer mismatch", err.Error()) + } + }) + + t.Run("issuer match ignoring trailing slash", func(t *testing.T) { + config := validClientCredentialsConfig() + // authServer.URL() has no trailing slash; configure with one to + // verify the comparison tolerates the difference (per RFC 8414 §3.3 + // normalization applied in oauthex.IssuersEqual). + config.Credentials.Issuer = authServer.URL() + "/" + handler, err := NewClientCredentialsHandler(config) + if err != nil { + t.Fatal(err) + } + + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: http.Header{}, + Body: http.NoBody, + } + req := httptest.NewRequest("GET", resourceURL, nil) + if err := handler.Authorize(t.Context(), req, resp); err != nil { + t.Fatalf("Authorize() unexpected error = %v", err) + } + }) + + t.Run("PRM via resource_metadata in challenge", func(t *testing.T) { + prmMux := http.NewServeMux() + prmMux.Handle("/custom-prm", auth.ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: resourceURL, + AuthorizationServers: []string{authServer.URL()}, + })) + prmServer := httptest.NewServer(prmMux) + t.Cleanup(prmServer.Close) + + handler, err := NewClientCredentialsHandler(validClientCredentialsConfig()) + if err != nil { + t.Fatal(err) + } + + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: http.Header{ + "WWW-Authenticate": []string{`Bearer resource_metadata="` + prmServer.URL + `/custom-prm"`}, + }, + Body: http.NoBody, + } + req := httptest.NewRequest("GET", resourceURL, nil) + if err := handler.Authorize(t.Context(), req, resp); err != nil { + t.Fatalf("Authorize with resource_metadata challenge failed: %v", err) + } + + ts, err := handler.TokenSource(t.Context()) + if err != nil { + t.Fatal(err) + } + if ts == nil { + t.Fatal("expected non-nil TokenSource") + } + }) +} + +func TestClientCredentialsHandler_ScopeAccumulation(t *testing.T) { + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + MetadataEndpointConfig: &oauthtest.MetadataEndpointConfig{ + ServeOAuthInsertedEndpoint: true, + }, + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test-client": {Secret: "test-secret"}, + }, + }, + ClientCredentialsConfig: &oauthtest.ClientCredentialsConfig{ + Enabled: true, + }, + TokenScopeFunc: func(requestedScope string) string { + // Simulate a server that never grants "write". + var granted []string + for _, s := range strings.Fields(requestedScope) { + if s != "write" { + granted = append(granted, s) + } + } + return strings.Join(granted, " ") + }, + }) + authServer.Start(t) + + resourceMux := http.NewServeMux() + resourceServer := httptest.NewServer(resourceMux) + t.Cleanup(resourceServer.Close) + resourceURL := resourceServer.URL + "/resource" + + resourceMux.Handle("/.well-known/oauth-protected-resource/resource", auth.ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: resourceURL, + AuthorizationServers: []string{authServer.URL()}, + })) + + handler, err := NewClientCredentialsHandler(validClientCredentialsConfig()) + if err != nil { + t.Fatal(err) + } + + // First authorization: 401 with scope="read write". + // The token response will only grant "read" (TokenScopeFunc strips "write"). + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + } + resp.Header.Set("WWW-Authenticate", `Bearer scope="read write"`) + req := httptest.NewRequest("GET", resourceURL, nil) + if err := handler.Authorize(t.Context(), req, resp); err != nil { + t.Fatalf("First Authorize failed: %v", err) + } + + // Verify only "read" was granted (the token omitted "write"). + issuer := authServer.URL() + if diff := cmp.Diff([]string{"read"}, handler.grantedScopes[issuer], cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" { + t.Errorf("After first Authorize, grantedScopes mismatch (-want +got):\n%s", diff) + } + + // Second authorization: 401 with scope="admin" (simulating step-up). + // Accumulated scopes should be "read" (previously granted) + "admin" (new), + // but NOT "write" (requested but never granted). + resp2 := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + } + resp2.Header.Set("WWW-Authenticate", `Bearer scope="admin"`) + req2 := httptest.NewRequest("GET", resourceURL, nil) + if err := handler.Authorize(t.Context(), req2, resp2); err != nil { + t.Fatalf("Second Authorize failed: %v", err) + } + + // Verify accumulated scopes: "read" (granted) + "admin" (new), no "write". + if diff := cmp.Diff([]string{"admin", "read"}, handler.grantedScopes[issuer], cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" { + t.Errorf("After second Authorize, grantedScopes mismatch (-want +got):\n%s", diff) + } +} + +func TestSelectTokenAuthMethod(t *testing.T) { + tests := []struct { + name string + supported []string + want oauth2.AuthStyle + }{ + { + name: "prefers client_secret_post", + supported: []string{"client_secret_basic", "client_secret_post"}, + want: oauth2.AuthStyleInParams, + }, + { + name: "falls back to client_secret_basic", + supported: []string{"client_secret_basic"}, + want: oauth2.AuthStyleInHeader, + }, + { + name: "auto detect when none supported", + supported: []string{"private_key_jwt"}, + want: oauth2.AuthStyleAutoDetect, + }, + { + name: "auto detect when empty", + supported: nil, + want: oauth2.AuthStyleAutoDetect, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := selectTokenAuthMethod(tc.supported) + if got != tc.want { + t.Errorf("selectTokenAuthMethod(%v) = %v, want %v", tc.supported, got, tc.want) + } + }) + } +} diff --git a/auth/extauth/enterprise_handler.go b/auth/extauth/enterprise_handler.go new file mode 100644 index 0000000..934af53 --- /dev/null +++ b/auth/extauth/enterprise_handler.go @@ -0,0 +1,254 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +// Package extauth provides OAuth handler implementations for MCP authorization extensions. +// This package implements Enterprise Managed Authorization as defined in SEP-990. + +package extauth + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" +) + +// grantTypeJWTBearer is the grant type for RFC 7523 JWT Bearer authorization grant. +const grantTypeJWTBearer = "urn:ietf:params:oauth:grant-type:jwt-bearer" + +// IDTokenFetcher is called to obtain an ID Token from the enterprise IdP. +// This is typically done via OIDC login flow where the user authenticates +// with their enterprise identity provider. +// +// Returns an oauth2.Token where Extra("id_token") contains the OpenID Connect ID Token (JWT). +type IDTokenFetcher func(ctx context.Context) (*oauth2.Token, error) + +// EnterpriseHandlerConfig is the configuration for [EnterpriseHandler]. +type EnterpriseHandlerConfig struct { + // IdP configuration (where the user authenticates) + + // IdPIssuerURL is the enterprise IdP's issuer URL (e.g., "https://acme.okta.com"). + // Used for OIDC discovery to find the token endpoint. + // REQUIRED. + IdPIssuerURL string + + // IdPCredentials contains the MCP Client's credentials registered at the IdP. + // REQUIRED. These credentials are used for token exchange at the IdP. + // The ClientID is always required. ClientSecretAuth is optional and only needed + // if the IdP requires client authentication (confidential clients). + IdPCredentials *oauthex.ClientCredentials + + // MCP Server configuration (the resource being accessed) + + // MCPAuthServerURL is the MCP Server's authorization server issuer URL. + // Used as the audience for token exchange and for metadata discovery. + // REQUIRED. + MCPAuthServerURL string + + // MCPResourceURI is the MCP Server's resource identifier (RFC 9728). + // Used as the resource parameter in token exchange. + // REQUIRED. + MCPResourceURI string + + // MCPCredentials contains the MCP Client's credentials registered at the MCP Server. + // REQUIRED. These credentials are used for JWT Bearer grant at the MCP Server. + // The ClientID is always required. ClientSecretAuth is optional and only needed + // if the MCP Server requires client authentication. + MCPCredentials *oauthex.ClientCredentials + + // MCPScopes is the list of scopes to request at the MCP Server. + // OPTIONAL. + MCPScopes []string + + // IDTokenFetcher is called to obtain an ID Token when authorization is needed. + // The implementation should handle the OIDC login flow (e.g., browser redirect, + // callback handling) and return the ID token. + // REQUIRED. + IDTokenFetcher IDTokenFetcher + + // HTTPClient is an optional HTTP client for customization. + // If nil, http.DefaultClient is used. + // OPTIONAL. + HTTPClient *http.Client +} + +// EnterpriseHandler is an implementation of [auth.OAuthHandler] that uses +// Enterprise Managed Authorization (SEP-990) to obtain access tokens. +// +// The flow consists of: +// 1. OIDC Login: User authenticates with enterprise IdP → ID Token +// 2. Token Exchange (RFC 8693): ID Token → ID-JAG at IdP +// 3. JWT Bearer Grant (RFC 7523): ID-JAG → Access Token at MCP Server +type EnterpriseHandler struct { + config *EnterpriseHandlerConfig + + // tokenSource is the token source obtained after authorization. + tokenSource oauth2.TokenSource +} + +// Compile-time check that EnterpriseHandler implements auth.OAuthHandler. +var _ auth.OAuthHandler = (*EnterpriseHandler)(nil) + +// NewEnterpriseHandler creates a new EnterpriseHandler. +// It performs validation of the configuration and returns an error if invalid. +func NewEnterpriseHandler(config *EnterpriseHandlerConfig) (*EnterpriseHandler, error) { + if config == nil { + return nil, errors.New("config must be provided") + } + if config.IdPIssuerURL == "" { + return nil, errors.New("IdPIssuerURL is required") + } + if config.IdPCredentials == nil { + return nil, errors.New("IdPCredentials is required") + } + if err := config.IdPCredentials.Validate(); err != nil { + return nil, fmt.Errorf("invalid IdPCredentials: %w", err) + } + if config.MCPAuthServerURL == "" { + return nil, errors.New("MCPAuthServerURL is required") + } + if config.MCPResourceURI == "" { + return nil, errors.New("MCPResourceURI is required") + } + if config.MCPCredentials == nil { + return nil, errors.New("MCPCredentials is required") + } + if err := config.MCPCredentials.Validate(); err != nil { + return nil, fmt.Errorf("invalid MCPCredentials: %w", err) + } + if config.IDTokenFetcher == nil { + return nil, errors.New("IDTokenFetcher is required") + } + return &EnterpriseHandler{config: config}, nil +} + +// TokenSource returns the token source for outgoing requests. +// Returns nil if authorization has not been performed yet. +func (h *EnterpriseHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) { + return h.tokenSource, nil +} + +// Authorize performs the Enterprise Managed Authorization flow. +// It is called when a request fails with 401 or 403. +func (h *EnterpriseHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error { + defer resp.Body.Close() + defer io.Copy(io.Discard, resp.Body) + + httpClient := h.config.HTTPClient + if httpClient == nil { + httpClient = http.DefaultClient + } + + // Step 1: Get ID Token via the configured fetcher (e.g., OIDC login) + oidcToken, err := h.config.IDTokenFetcher(ctx) + if err != nil { + return fmt.Errorf("failed to obtain ID token: %w", err) + } + + // Extract ID token from the oauth2.Token + idToken, ok := oidcToken.Extra("id_token").(string) + if !ok || idToken == "" { + return fmt.Errorf("id_token not found in OIDC token response") + } + + // Step 2: Discover IdP token endpoint via OIDC discovery + idpMeta, err := auth.GetAuthServerMetadata(ctx, h.config.IdPIssuerURL, httpClient) + if err != nil { + return fmt.Errorf("failed to discover IdP metadata: %w", err) + } + if idpMeta == nil { + return fmt.Errorf("no authorization server metadata found for IdP: %s", h.config.IdPIssuerURL) + } + + // Step 3: Token Exchange (ID Token → ID-JAG) + tokenExchangeReq := &oauthex.TokenExchangeRequest{ + RequestedTokenType: oauthex.TokenTypeIDJAG, + Audience: h.config.MCPAuthServerURL, + Resource: h.config.MCPResourceURI, + Scope: h.config.MCPScopes, + SubjectToken: idToken, + SubjectTokenType: oauthex.TokenTypeIDToken, + } + + idJAGToken, err := oauthex.ExchangeToken( + ctx, + idpMeta.TokenEndpoint, + tokenExchangeReq, + h.config.IdPCredentials, + httpClient, + ) + if err != nil { + return fmt.Errorf("token exchange failed: %w", err) + } + + // Step 4: Discover MCP Server token endpoint + mcpMeta, err := auth.GetAuthServerMetadata(ctx, h.config.MCPAuthServerURL, httpClient) + if err != nil { + return fmt.Errorf("failed to discover MCP auth server metadata: %w", err) + } + if mcpMeta == nil { + return fmt.Errorf("no authorization server metadata found for MCP server: %s", h.config.MCPAuthServerURL) + } + + // Step 5: JWT Bearer Grant (ID-JAG → Access Token) + // The ID-JAG is in the AccessToken field of the token (despite the name) + accessToken, err := exchangeJWTBearer( + ctx, + mcpMeta.TokenEndpoint, + idJAGToken.AccessToken, + h.config.MCPCredentials, + httpClient, + ) + if err != nil { + return fmt.Errorf("JWT bearer grant failed: %w", err) + } + + // Store the token source for subsequent requests + h.tokenSource = oauth2.StaticTokenSource(accessToken) + return nil +} + +// exchangeJWTBearer exchanges an Identity Assertion JWT Authorization Grant (ID-JAG) +// for an access token using JWT Bearer Grant per RFC 7523. +func exchangeJWTBearer( + ctx context.Context, + tokenEndpoint string, + assertion string, + clientCreds *oauthex.ClientCredentials, + httpClient *http.Client, +) (*oauth2.Token, error) { + cfg := &oauth2.Config{ + ClientID: clientCreds.ClientID, + Endpoint: oauth2.Endpoint{ + TokenURL: tokenEndpoint, + AuthStyle: oauth2.AuthStyleInParams, + }, + } + // Set ClientSecret if ClientSecretAuth is configured + if clientCreds.ClientSecretAuth != nil { + cfg.ClientSecret = clientCreds.ClientSecretAuth.ClientSecret + } + + if httpClient == nil { + httpClient = http.DefaultClient + } + ctxWithClient := context.WithValue(ctx, oauth2.HTTPClient, httpClient) + + token, err := cfg.Exchange( + ctxWithClient, + "", + oauth2.SetAuthURLParam("grant_type", grantTypeJWTBearer), + oauth2.SetAuthURLParam("assertion", assertion), + ) + if err != nil { + return nil, fmt.Errorf("JWT bearer grant request failed: %w", err) + } + + return token, nil +} diff --git a/auth/extauth/enterprise_handler_test.go b/auth/extauth/enterprise_handler_test.go new file mode 100644 index 0000000..86f0414 --- /dev/null +++ b/auth/extauth/enterprise_handler_test.go @@ -0,0 +1,279 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package extauth + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/modelcontextprotocol/go-sdk/internal/oauthtest" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" +) + +func validEnterpriseHandlerConfig() *EnterpriseHandlerConfig { + return &EnterpriseHandlerConfig{ + IdPIssuerURL: "https://idp.example.com", + IdPCredentials: &oauthex.ClientCredentials{ClientID: "idp_client_id"}, + MCPAuthServerURL: "https://mcp-auth.example.com", + MCPResourceURI: "https://mcp.example.com", + MCPCredentials: &oauthex.ClientCredentials{ClientID: "mcp_client_id"}, + IDTokenFetcher: func(ctx context.Context) (*oauth2.Token, error) { + token := &oauth2.Token{AccessToken: "mock", TokenType: "Bearer"} + return token.WithExtra(map[string]any{"id_token": "mock_id_token"}), nil + }, + } +} + +func TestNewEnterpriseHandler_Validation(t *testing.T) { + tests := []struct { + name string + config *EnterpriseHandlerConfig + wantError string + }{ + { + name: "nil config", + config: nil, + wantError: "config must be provided", + }, + { + name: "missing IdPIssuerURL", + config: func() *EnterpriseHandlerConfig { + c := validEnterpriseHandlerConfig() + c.IdPIssuerURL = "" + return c + }(), + wantError: "IdPIssuerURL is required", + }, + { + name: "nil IdPCredentials", + config: func() *EnterpriseHandlerConfig { + c := validEnterpriseHandlerConfig() + c.IdPCredentials = nil + return c + }(), + wantError: "IdPCredentials is required", + }, + { + name: "invalid IdPCredentials - empty ClientID", + config: func() *EnterpriseHandlerConfig { + c := validEnterpriseHandlerConfig() + c.IdPCredentials = &oauthex.ClientCredentials{ClientID: ""} + return c + }(), + wantError: "invalid IdPCredentials", + }, + { + name: "invalid IdPCredentials - empty ClientSecret in ClientSecretAuth", + config: func() *EnterpriseHandlerConfig { + c := validEnterpriseHandlerConfig() + c.IdPCredentials = &oauthex.ClientCredentials{ + ClientID: "idp_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ClientSecret: ""}, + } + return c + }(), + wantError: "invalid IdPCredentials", + }, + { + name: "missing MCPAuthServerURL", + config: func() *EnterpriseHandlerConfig { + c := validEnterpriseHandlerConfig() + c.MCPAuthServerURL = "" + return c + }(), + wantError: "MCPAuthServerURL is required", + }, + { + name: "missing MCPResourceURI", + config: func() *EnterpriseHandlerConfig { + c := validEnterpriseHandlerConfig() + c.MCPResourceURI = "" + return c + }(), + wantError: "MCPResourceURI is required", + }, + { + name: "nil MCPCredentials", + config: func() *EnterpriseHandlerConfig { + c := validEnterpriseHandlerConfig() + c.MCPCredentials = nil + return c + }(), + wantError: "MCPCredentials is required", + }, + { + name: "invalid MCPCredentials - empty ClientID", + config: func() *EnterpriseHandlerConfig { + c := validEnterpriseHandlerConfig() + c.MCPCredentials = &oauthex.ClientCredentials{ClientID: ""} + return c + }(), + wantError: "invalid MCPCredentials", + }, + { + name: "missing IDTokenFetcher", + config: func() *EnterpriseHandlerConfig { + c := validEnterpriseHandlerConfig() + c.IDTokenFetcher = nil + return c + }(), + wantError: "IDTokenFetcher is required", + }, + { + name: "valid public clients", + config: validEnterpriseHandlerConfig(), + wantError: "", + }, + { + name: "valid confidential clients", + config: func() *EnterpriseHandlerConfig { + c := validEnterpriseHandlerConfig() + c.IdPCredentials.ClientSecretAuth = &oauthex.ClientSecretAuth{ClientSecret: "idp_secret"} + c.MCPCredentials.ClientSecretAuth = &oauthex.ClientSecretAuth{ClientSecret: "mcp_secret"} + return c + }(), + wantError: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + handler, err := NewEnterpriseHandler(tt.config) + if tt.wantError != "" { + if err == nil { + t.Fatalf("expected error containing %q, got nil", tt.wantError) + } + if !strings.Contains(err.Error(), tt.wantError) { + t.Fatalf("expected error containing %q, got %v", tt.wantError, err) + } + } else { + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if handler == nil { + t.Fatal("expected handler to be non-nil") + } + } + }) + } +} + +func TestEnterpriseHandler_Authorize_E2E(t *testing.T) { + const idJAGToken = "id-jag-token-from-idp" + + idpServer := oauthtest.NewFakeIdPServer(oauthtest.IdPConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "idp_client_id": {Secret: "idp_secret"}, + }, + TokenExchangeConfig: &oauthtest.TokenExchangeConfig{ + IDJAGToken: idJAGToken, + }, + }) + idpServer.Start(t) + + mcpAuthServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "mcp_client_id": {Secret: "mcp_secret"}, + }, + }, + JWTBearerConfig: &oauthtest.JWTBearerConfig{ + ValidAssertions: []string{idJAGToken}, + }, + }) + mcpAuthServer.Start(t) + + config := validEnterpriseHandlerConfig() + config.IdPIssuerURL = idpServer.URL() + config.IdPCredentials.ClientSecretAuth = &oauthex.ClientSecretAuth{ClientSecret: "idp_secret"} + config.MCPAuthServerURL = mcpAuthServer.URL() + config.MCPCredentials.ClientSecretAuth = &oauthex.ClientSecretAuth{ClientSecret: "mcp_secret"} + config.MCPScopes = []string{"read", "write"} + config.IDTokenFetcher = func(ctx context.Context) (*oauth2.Token, error) { + token := &oauth2.Token{AccessToken: "mock_access_token", TokenType: "Bearer"} + return token.WithExtra(map[string]any{"id_token": "mock_id_token_from_user_login"}), nil + } + + handler, err := NewEnterpriseHandler(config) + if err != nil { + t.Fatalf("NewEnterpriseHandler failed: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "https://mcp.example.com/api", nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + + if err := handler.Authorize(context.Background(), req, resp); err != nil { + t.Fatalf("Authorize failed: %v", err) + } + + tokenSource, err := handler.TokenSource(context.Background()) + if err != nil { + t.Fatalf("TokenSource failed: %v", err) + } + if tokenSource == nil { + t.Fatal("expected token source to be set after authorization") + } + + token, err := tokenSource.Token() + if err != nil { + t.Fatalf("Token() failed: %v", err) + } + if token.AccessToken != "test_access_token" { + t.Errorf("AccessToken = %q, want %q", token.AccessToken, "test_access_token") + } +} + +func TestEnterpriseHandler_Authorize_IDTokenFetcherError(t *testing.T) { + config := validEnterpriseHandlerConfig() + config.IDTokenFetcher = func(ctx context.Context) (*oauth2.Token, error) { + return nil, fmt.Errorf("user cancelled login") + } + + handler, err := NewEnterpriseHandler(config) + if err != nil { + t.Fatalf("NewEnterpriseHandler failed: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "https://mcp.example.com/api", nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + + err = handler.Authorize(context.Background(), req, resp) + if err == nil { + t.Fatal("expected error from Authorize, got nil") + } + if !strings.Contains(err.Error(), "failed to obtain ID token") { + t.Errorf("error = %v, want error containing %q", err, "failed to obtain ID token") + } +} + +func TestEnterpriseHandler_TokenSource_BeforeAuthorization(t *testing.T) { + handler, err := NewEnterpriseHandler(validEnterpriseHandlerConfig()) + if err != nil { + t.Fatalf("NewEnterpriseHandler failed: %v", err) + } + + tokenSource, err := handler.TokenSource(context.Background()) + if err != nil { + t.Fatalf("TokenSource failed: %v", err) + } + if tokenSource != nil { + t.Errorf("expected nil token source before authorization, got %v", tokenSource) + } +} diff --git a/auth/extauth/oidc_login.go b/auth/extauth/oidc_login.go new file mode 100644 index 0000000..d9d259b --- /dev/null +++ b/auth/extauth/oidc_login.go @@ -0,0 +1,240 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +// This file implements OIDC Authorization Code flow for obtaining ID tokens +// as part of Enterprise Managed Authorization (SEP-990). +// See https://openid.net/specs/openid-connect-core-1_0.html + +package extauth + +import ( + "context" + "crypto/rand" + "fmt" + "net/http" + "net/url" + "slices" + "strings" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" +) + +// OIDCLoginConfig configures the OIDC Authorization Code flow for obtaining +// an ID Token. This is used with [PerformOIDCLogin] to authenticate users +// with an enterprise IdP before calling the Enterprise Managed Authorization flow. +type OIDCLoginConfig struct { + // IssuerURL is the IdP's issuer URL (e.g., "https://acme.okta.com"). + // REQUIRED. + IssuerURL string + // Credentials contains the MCP Client's credentials registered at the IdP. + // The ClientID field is REQUIRED. The ClientSecret field is OPTIONAL + // (only required if the client is confidential, not a public client). + // REQUIRED (struct itself), but ClientSecret field can be empty. + Credentials *oauthex.ClientCredentials + // RedirectURL is the OAuth2 redirect URI registered with the IdP. + // This must match exactly what was registered with the IdP. + // REQUIRED. + RedirectURL string + // Scopes are the OAuth2/OIDC scopes to request. + // "openid" is REQUIRED for OIDC. Common values: ["openid", "profile", "email"] + // REQUIRED. + Scopes []string + // LoginHint is a hint to the IdP about the user's identity. + // Some IdPs may require this (e.g., as an email address for routing to SSO providers). + // Example: "user@example.com" + // OPTIONAL. + LoginHint string + // HTTPClient is the HTTP client for making requests. + // If nil, http.DefaultClient is used. + // OPTIONAL. + HTTPClient *http.Client +} + +// PerformOIDCLogin performs the complete OIDC Authorization Code flow with PKCE +// in a single function call. This is the recommended approach for obtaining an +// ID Token for use with [EnterpriseHandler]. +// +// Returns an oauth2.Token where: +// - Extra("id_token") contains the OpenID Connect ID Token (JWT) +// - AccessToken contains the OAuth2 access token (if issued by IdP) +// - RefreshToken contains the OAuth2 refresh token (if issued by IdP) +// - TokenType is the token type (typically "Bearer") +// - Expiry is when the token expires +func PerformOIDCLogin( + ctx context.Context, + config *OIDCLoginConfig, + authCodeFetcher auth.AuthorizationCodeFetcher, +) (*oauth2.Token, error) { + if authCodeFetcher == nil { + return nil, fmt.Errorf("authCodeFetcher is required") + } + + authReq, oauth2Config, err := initiateOIDCLogin(ctx, config) + if err != nil { + return nil, fmt.Errorf("failed to initiate OIDC login: %w", err) + } + + authResult, err := authCodeFetcher(ctx, &auth.AuthorizationArgs{URL: authReq.authURL}) + if err != nil { + return nil, fmt.Errorf("failed to fetch authorization code: %w", err) + } + + if authResult.State != authReq.state { + return nil, fmt.Errorf("state mismatch: expected %q, got %q", authReq.state, authResult.State) + } + + tokens, err := completeOIDCLogin(ctx, config, oauth2Config, authResult.Code, authReq.codeVerifier) + if err != nil { + return nil, fmt.Errorf("failed to complete OIDC login: %w", err) + } + + return tokens, nil +} + +// oidcAuthorizationRequest holds internal state for OIDC authorization. +type oidcAuthorizationRequest struct { + authURL string + state string + codeVerifier string +} + +// initiateOIDCLogin initiates an OIDC Authorization Code flow with PKCE. +func initiateOIDCLogin( + ctx context.Context, + config *OIDCLoginConfig, +) (*oidcAuthorizationRequest, *oauth2.Config, error) { + if config == nil { + return nil, nil, fmt.Errorf("config is required") + } + if config.IssuerURL == "" { + return nil, nil, fmt.Errorf("IssuerURL is required") + } + if config.Credentials == nil || config.Credentials.ClientID == "" { + return nil, nil, fmt.Errorf("Credentials.ClientID is required") + } + if config.RedirectURL == "" { + return nil, nil, fmt.Errorf("RedirectURL is required") + } + if len(config.Scopes) == 0 { + return nil, nil, fmt.Errorf("at least one scope is required (must include 'openid')") + } + + if !slices.Contains(config.Scopes, "openid") { + return nil, nil, fmt.Errorf("the 'openid' scope is required for OIDC") + } + + if err := checkURLScheme(config.IssuerURL); err != nil { + return nil, nil, fmt.Errorf("invalid IssuerURL: %w", err) + } + if err := checkURLScheme(config.RedirectURL); err != nil { + return nil, nil, fmt.Errorf("invalid RedirectURL: %w", err) + } + + httpClient := config.HTTPClient + if httpClient == nil { + httpClient = http.DefaultClient + } + meta, err := auth.GetAuthServerMetadata(ctx, config.IssuerURL, httpClient) + if err != nil { + return nil, nil, fmt.Errorf("failed to discover OIDC metadata: %w", err) + } + if meta == nil { + return nil, nil, fmt.Errorf("no authorization server metadata found for OIDC issuer: %s", config.IssuerURL) + } + if meta.AuthorizationEndpoint == "" { + return nil, nil, fmt.Errorf("authorization_endpoint not found in OIDC metadata") + } + + codeVerifier := oauth2.GenerateVerifier() + state := rand.Text() + + oauth2Config := &oauth2.Config{ + ClientID: config.Credentials.ClientID, + RedirectURL: config.RedirectURL, + Scopes: config.Scopes, + Endpoint: oauth2.Endpoint{ + AuthURL: meta.AuthorizationEndpoint, + TokenURL: meta.TokenEndpoint, + }, + } + // Set ClientSecret if ClientSecretAuth is configured + if config.Credentials.ClientSecretAuth != nil { + oauth2Config.ClientSecret = config.Credentials.ClientSecretAuth.ClientSecret + } + + authURLOpts := []oauth2.AuthCodeOption{ + oauth2.S256ChallengeOption(codeVerifier), + } + if config.LoginHint != "" { + authURLOpts = append(authURLOpts, oauth2.SetAuthURLParam("login_hint", config.LoginHint)) + } + authURL := oauth2Config.AuthCodeURL(state, authURLOpts...) + + return &oidcAuthorizationRequest{ + authURL: authURL, + state: state, + codeVerifier: codeVerifier, + }, oauth2Config, nil +} + +// completeOIDCLogin completes the OIDC Authorization Code flow by exchanging +// the authorization code for tokens. +func completeOIDCLogin( + ctx context.Context, + config *OIDCLoginConfig, + oauth2Config *oauth2.Config, + authCode string, + codeVerifier string, +) (*oauth2.Token, error) { + if authCode == "" { + return nil, fmt.Errorf("authCode is required") + } + if codeVerifier == "" { + return nil, fmt.Errorf("codeVerifier is required") + } + + httpClient := config.HTTPClient + if httpClient == nil { + httpClient = http.DefaultClient + } + ctxWithClient := context.WithValue(ctx, oauth2.HTTPClient, httpClient) + + oauth2Token, err := oauth2Config.Exchange( + ctxWithClient, + authCode, + oauth2.VerifierOption(codeVerifier), + ) + if err != nil { + return nil, fmt.Errorf("token exchange failed: %w", err) + } + + // Validate that id_token is present in the response + idToken, ok := oauth2Token.Extra("id_token").(string) + if !ok || idToken == "" { + return nil, fmt.Errorf("id_token not found in token response") + } + + return oauth2Token, nil +} + +// checkURLScheme ensures that its argument is a valid URL with a scheme +// that prevents XSS attacks. +// See #526. +// Note: a copy of this function exists in oauthex/oauth2.go; keep these in sync. +func checkURLScheme(u string) error { + if u == "" { + return nil + } + uu, err := url.Parse(u) + if err != nil { + return err + } + scheme := strings.ToLower(uu.Scheme) + if scheme == "javascript" || scheme == "data" || scheme == "vbscript" { + return fmt.Errorf("URL has disallowed scheme %q", scheme) + } + return nil +} diff --git a/auth/extauth/oidc_login_test.go b/auth/extauth/oidc_login_test.go new file mode 100644 index 0000000..9cd9b93 --- /dev/null +++ b/auth/extauth/oidc_login_test.go @@ -0,0 +1,263 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package extauth + +import ( + "context" + "fmt" + "net/http" + "net/url" + "strings" + "testing" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/internal/oauthtest" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" +) + +const testRedirectURL = "http://localhost:18927/callback" + +func validOIDCLoginConfig(issuerURL string) *OIDCLoginConfig { + return &OIDCLoginConfig{ + IssuerURL: issuerURL, + Credentials: &oauthex.ClientCredentials{ + ClientID: "test-client", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "test-secret", + }, + }, + RedirectURL: testRedirectURL, + Scopes: []string{"openid", "profile", "email"}, + } +} + +func TestPerformOIDCLogin(t *testing.T) { + idpServer := oauthtest.NewFakeIdPServer(oauthtest.IdPConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test-client": { + Secret: "test-secret", + RedirectURIs: []string{testRedirectURL}, + }, + }, + }) + idpServer.Start(t) + + config := validOIDCLoginConfig(idpServer.URL()) + + // fetchAuthCode visits the authorization URL on the fake IdP server, + // follows the redirect, and extracts the authorization code and state. + fetchAuthCode := func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + client := &http.Client{ + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } + resp, err := client.Get(args.URL) + if err != nil { + return nil, fmt.Errorf("failed to visit auth URL: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusFound { + return nil, fmt.Errorf("expected redirect, got status %d", resp.StatusCode) + } + loc, err := resp.Location() + if err != nil { + return nil, fmt.Errorf("missing Location header: %w", err) + } + return &auth.AuthorizationResult{ + Code: loc.Query().Get("code"), + State: loc.Query().Get("state"), + }, nil + } + + verifyOIDCToken := func(token *oauth2.Token) error { + idToken, ok := token.Extra("id_token").(string) + if !ok || idToken == "" { + return fmt.Errorf("id_token is missing or empty") + } + parts := strings.Split(idToken, ".") + if len(parts) != 3 { + return fmt.Errorf("id_token is not a JWT (expected 3 parts, got %d)", len(parts)) + } + if token.AccessToken == "" { + return fmt.Errorf("access token is empty") + } + if token.TokenType != "Bearer" { + return fmt.Errorf("token type = %q, want %q", token.TokenType, "Bearer") + } + return nil + } + + t.Run("successful flow", func(t *testing.T) { + token, err := PerformOIDCLogin(context.Background(), config, fetchAuthCode) + if err != nil { + t.Fatalf("PerformOIDCLogin() error = %v", err) + } + if err := verifyOIDCToken(token); err != nil { + t.Errorf("invalid token returned: %v", err) + } + }) + + t.Run("with login_hint", func(t *testing.T) { + configWithHint := *config + configWithHint.LoginHint = "user@example.com" + + token, err := PerformOIDCLogin(context.Background(), &configWithHint, + func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + u, err := url.Parse(args.URL) + if err != nil { + return nil, fmt.Errorf("invalid authURL: %w", err) + } + if got := u.Query().Get("login_hint"); got != "user@example.com" { + t.Errorf("login_hint = %q, want %q", got, "user@example.com") + } + return fetchAuthCode(ctx, args) + }) + if err != nil { + t.Fatalf("PerformOIDCLogin() error = %v", err) + } + if err := verifyOIDCToken(token); err != nil { + t.Errorf("invalid token returned: %v", err) + } + }) + + t.Run("without login_hint", func(t *testing.T) { + token, err := PerformOIDCLogin(context.Background(), config, + func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + u, err := url.Parse(args.URL) + if err != nil { + return nil, fmt.Errorf("invalid authURL: %w", err) + } + if u.Query().Has("login_hint") { + t.Errorf("login_hint should be absent, got %q", u.Query().Get("login_hint")) + } + return fetchAuthCode(ctx, args) + }) + if err != nil { + t.Fatalf("PerformOIDCLogin() error = %v", err) + } + if err := verifyOIDCToken(token); err != nil { + t.Errorf("invalid token returned: %v", err) + } + }) + + t.Run("state mismatch", func(t *testing.T) { + _, err := PerformOIDCLogin(context.Background(), config, + func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + return &auth.AuthorizationResult{ + Code: "mock-auth-code", + State: "wrong-state", + }, nil + }) + + if err == nil { + t.Fatal("expected error for state mismatch, got nil") + } + if !strings.Contains(err.Error(), "state mismatch") { + t.Errorf("error = %v, want error containing %q", err, "state mismatch") + } + }) + + t.Run("fetcher error", func(t *testing.T) { + _, err := PerformOIDCLogin(context.Background(), config, + func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + return nil, fmt.Errorf("user cancelled") + }) + + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "user cancelled") { + t.Errorf("error = %v, want error containing %q", err, "user cancelled") + } + }) +} + +func TestPerformOIDCLogin_ValidationErrors(t *testing.T) { + const nonexistentIssuer = "https://idp.example.com" + + noopFetcher := func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + return nil, fmt.Errorf("should not be called") + } + + tests := []struct { + name string + config *OIDCLoginConfig + fetcher auth.AuthorizationCodeFetcher + wantErr string + }{ + { + name: "nil config", + config: nil, + fetcher: noopFetcher, + wantErr: "config is required", + }, + { + name: "nil fetcher", + config: validOIDCLoginConfig(nonexistentIssuer), + fetcher: nil, + wantErr: "authCodeFetcher is required", + }, + { + name: "missing IssuerURL", + config: validOIDCLoginConfig(""), + fetcher: noopFetcher, + wantErr: "IssuerURL is required", + }, + { + name: "missing ClientID", + config: func() *OIDCLoginConfig { + c := validOIDCLoginConfig(nonexistentIssuer) + c.Credentials = &oauthex.ClientCredentials{} + return c + }(), + fetcher: noopFetcher, + wantErr: "ClientID is required", + }, + { + name: "missing RedirectURL", + config: func() *OIDCLoginConfig { + c := validOIDCLoginConfig(nonexistentIssuer) + c.RedirectURL = "" + return c + }(), + fetcher: noopFetcher, + wantErr: "RedirectURL is required", + }, + { + name: "missing Scopes", + config: func() *OIDCLoginConfig { + c := validOIDCLoginConfig(nonexistentIssuer) + c.Scopes = nil + return c + }(), + fetcher: noopFetcher, + wantErr: "at least one scope is required", + }, + { + name: "missing openid scope", + config: func() *OIDCLoginConfig { + c := validOIDCLoginConfig(nonexistentIssuer) + c.Scopes = []string{"profile", "email"} + return c + }(), + fetcher: noopFetcher, + wantErr: "openid", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := PerformOIDCLogin(context.Background(), tt.config, tt.fetcher) + if err == nil { + t.Fatalf("expected error containing %q, got nil", tt.wantErr) + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Errorf("error = %v, want error containing %q", err, tt.wantErr) + } + }) + } +} diff --git a/auth/shared.go b/auth/shared.go new file mode 100644 index 0000000..e58d438 --- /dev/null +++ b/auth/shared.go @@ -0,0 +1,73 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +// This file contains shared utilities for OAuth handlers. + +package auth + +import ( + "context" + "net/http" + "net/url" + "strings" + + "github.com/modelcontextprotocol/go-sdk/oauthex" +) + +// GetAuthServerMetadata fetches authorization server metadata for the given issuer URL. +// It tries standard well-known endpoints (OAuth 2.0 and OIDC) and returns the first successful result. +// +// Returns (nil, nil) when no metadata endpoints respond (404s), allowing callers to implement +// fallback logic. Returns an error for any non-client error (network failures, invalid JSON, etc.). +func GetAuthServerMetadata(ctx context.Context, issuerURL string, httpClient *http.Client) (*oauthex.AuthServerMeta, error) { + for _, metadataURL := range authorizationServerMetadataURLs(issuerURL) { + asm, err := oauthex.GetAuthServerMeta(ctx, metadataURL, issuerURL, httpClient) + if err != nil { + return nil, err + } + if asm != nil { + return asm, nil + } + } + return nil, nil +} + +// authorizationServerMetadataURLs returns a list of URLs to try when looking for +// authorization server metadata as mandated by the MCP specification: +// https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#authorization-server-metadata-discovery. +func authorizationServerMetadataURLs(issuerURL string) []string { + var urls []string + + baseURL, err := url.Parse(issuerURL) + if err != nil { + return nil + } + + // RFC 8414, section 3.1 requires any terminating "/" to be removed from the + // issuer identifier before the well-known suffix is inserted, so an issuer + // such as "https://auth.example.com/" has no path component at all. + issuerPath := strings.Trim(baseURL.Path, "/") + + if issuerPath == "" { + // "OAuth 2.0 Authorization Server Metadata". + baseURL.Path = "/.well-known/oauth-authorization-server" + urls = append(urls, baseURL.String()) + // "OpenID Connect Discovery 1.0". + baseURL.Path = "/.well-known/openid-configuration" + urls = append(urls, baseURL.String()) + return urls + } + + // "OAuth 2.0 Authorization Server Metadata with path insertion". + baseURL.Path = "/.well-known/oauth-authorization-server/" + issuerPath + urls = append(urls, baseURL.String()) + // "OpenID Connect Discovery 1.0 with path insertion". + baseURL.Path = "/.well-known/openid-configuration/" + issuerPath + urls = append(urls, baseURL.String()) + // "OpenID Connect Discovery 1.0 with path appending". + baseURL.Path = "/" + issuerPath + "/.well-known/openid-configuration" + urls = append(urls, baseURL.String()) + + return urls +} diff --git a/auth/shared_test.go b/auth/shared_test.go new file mode 100644 index 0000000..34e0a35 --- /dev/null +++ b/auth/shared_test.go @@ -0,0 +1,169 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package auth + +import ( + "net/http" + "slices" + "testing" + + "github.com/modelcontextprotocol/go-sdk/internal/oauthtest" +) + +func TestAuthorizationServerMetadataURLs(t *testing.T) { + tests := []struct { + issuerURL string + want []string + }{ + { + issuerURL: "https://auth.example.com", + want: []string{ + "https://auth.example.com/.well-known/oauth-authorization-server", + "https://auth.example.com/.well-known/openid-configuration", + }, + }, + { + issuerURL: "https://auth.example.com/", + want: []string{ + "https://auth.example.com/.well-known/oauth-authorization-server", + "https://auth.example.com/.well-known/openid-configuration", + }, + }, + { + issuerURL: "https://auth.example.com/tenant1", + want: []string{ + "https://auth.example.com/.well-known/oauth-authorization-server/tenant1", + "https://auth.example.com/.well-known/openid-configuration/tenant1", + "https://auth.example.com/tenant1/.well-known/openid-configuration", + }, + }, + { + issuerURL: "https://auth.example.com/tenant1/", + want: []string{ + "https://auth.example.com/.well-known/oauth-authorization-server/tenant1", + "https://auth.example.com/.well-known/openid-configuration/tenant1", + "https://auth.example.com/tenant1/.well-known/openid-configuration", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.issuerURL, func(t *testing.T) { + if got := authorizationServerMetadataURLs(tt.issuerURL); !slices.Equal(got, tt.want) { + t.Errorf("authorizationServerMetadataURLs(%q) =\n%v\nwant\n%v", tt.issuerURL, got, tt.want) + } + }) + } +} + +// TestGetAuthServerMetadataTrailingSlashIssuer checks that an issuer identifier +// written with a terminating slash still finds metadata served at the root +// well-known path. +func TestGetAuthServerMetadataTrailingSlashIssuer(t *testing.T) { + s := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + MetadataEndpointConfig: &oauthtest.MetadataEndpointConfig{ + ServeOAuthInsertedEndpoint: true, + }, + }) + s.Start(t) + + got, err := GetAuthServerMetadata(t.Context(), s.URL()+"/", http.DefaultClient) + if err != nil { + t.Fatalf("GetAuthServerMetadata() error = %v, want nil", err) + } + if got == nil { + t.Fatal("GetAuthServerMetadata() got nil, want metadata") + } + if got.Issuer != s.URL() { + t.Errorf("GetAuthServerMetadata() issuer = %q, want %q", got.Issuer, s.URL()) + } +} + +func TestGetAuthServerMetadata(t *testing.T) { + tests := []struct { + name string + issuerPath string + endpointConfig *oauthtest.MetadataEndpointConfig + wantNil bool + }{ + { + name: "OAuthEndpoint_Root", + issuerPath: "", + endpointConfig: &oauthtest.MetadataEndpointConfig{ + ServeOAuthInsertedEndpoint: true, + }, + }, + { + name: "OpenIDEndpoint_Root", + issuerPath: "", + endpointConfig: &oauthtest.MetadataEndpointConfig{ + ServeOpenIDInsertedEndpoint: true, + }, + }, + { + name: "OAuthEndpoint_Path", + issuerPath: "/oauth", + endpointConfig: &oauthtest.MetadataEndpointConfig{ + ServeOAuthInsertedEndpoint: true, + }, + }, + { + name: "OpenIDEndpoint_Path", + issuerPath: "/openid", + endpointConfig: &oauthtest.MetadataEndpointConfig{ + ServeOpenIDInsertedEndpoint: true, + }, + }, + { + name: "OpenIDAppendedEndpoint_Path", + issuerPath: "/openid", + endpointConfig: &oauthtest.MetadataEndpointConfig{ + ServeOpenIDAppendedEndpoint: true, + }, + }, + { + name: "NoMetadata", + issuerPath: "", + endpointConfig: &oauthtest.MetadataEndpointConfig{ + // All metadata endpoints disabled. + ServeOAuthInsertedEndpoint: false, + ServeOpenIDInsertedEndpoint: false, + }, + wantNil: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + IssuerPath: tt.issuerPath, + MetadataEndpointConfig: tt.endpointConfig, + }) + s.Start(t) + issuerURL := s.URL() + tt.issuerPath + + got, err := GetAuthServerMetadata(t.Context(), issuerURL, http.DefaultClient) + if tt.wantNil { + // When no metadata is found, GetAuthServerMetadata returns (nil, nil). + if err != nil { + t.Fatalf("GetAuthServerMetadata() unexpected error = %v, want nil", err) + } + if got != nil { + t.Fatal("GetAuthServerMetadata() expected nil for no metadata, got metadata") + } + return + } + if err != nil { + t.Fatalf("GetAuthServerMetadata() error = %v, want nil", err) + } + if got == nil { + t.Fatal("GetAuthServerMetadata() got nil, want metadata") + } + if got.Issuer != issuerURL { + t.Errorf("GetAuthServerMetadata() issuer = %q, want %q", got.Issuer, issuerURL) + } + }) + } +} diff --git a/conformance/baseline.yml b/conformance/baseline.yml new file mode 100644 index 0000000..4867fe9 --- /dev/null +++ b/conformance/baseline.yml @@ -0,0 +1,14 @@ +client: + # OAuth mechanisms not implemented by the Go MCP SDK. + - auth/client-credentials-basic + - auth/client-credentials-jwt + - auth/enterprise-managed-authorization + + # SEP-1932 (DPoP Sender-Constrained Tokens) client-side support not yet + # implemented by the Go MCP SDK. + - auth/dpop + - auth/dpop-nonce + + # SEP-1933 (Workload Identity Federation JWT-bearer) client-side support + # not yet implemented by the Go MCP SDK. + - auth/wif-jwt-bearer diff --git a/conformance/everything-client/main.go b/conformance/everything-client/main.go new file mode 100644 index 0000000..b7d8ba0 --- /dev/null +++ b/conformance/everything-client/main.go @@ -0,0 +1,632 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +// The conformance client implements features required for MCP conformance testing. +// It mirrors the functionality of the TypeScript conformance client at +// https://github.com/modelcontextprotocol/typescript-sdk/blob/main/src/conformance/everything-client.ts +package main + +import ( + "context" + "encoding/json" + "fmt" + "log" + "net/http" + "net/url" + "os" + "slices" + "sort" + "strings" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/oauthex" +) + +// scenarioHandler is the function signature for all conformance test scenarios. +// It takes a context and the server URL to connect to. +type scenarioHandler func(ctx context.Context, serverURL string, configCtx map[string]any) error + +var ( + // registry stores all registered scenario handlers. + registry = make(map[string]scenarioHandler) +) + +// registerScenario registers a new scenario handler with the given name. +// This function should be called during init() by scenario implementations. +func registerScenario(name string, handler scenarioHandler) { + if _, exists := registry[name]; exists { + log.Fatalf("Scenario %q is already registered", name) + } + registry[name] = handler +} + +func init() { + registerScenario("initialize", runBasicClient) + registerScenario("tools_call", runToolsCallClient) + registerScenario("request-metadata", runRequestMetadataClient) + registerScenario("elicitation-sep1034-client-defaults", runElicitationDefaultsClient) + registerScenario("sse-retry", runSSERetryClient) + registerScenario("json-schema-ref-no-deref", runJSONSchemaRefNoDerefClient) + registerScenario("json-schema-2020-12-preservation", runJSONSchemaPreservationClient) + registerScenario("sep-2322-client-request-state", runMrtrClient) + registerScenario("http-standard-headers", runHTTPStandardHeadersClient) + registerScenario("http-custom-headers", runHTTPCustomHeadersClient) + registerScenario("http-invalid-tool-headers", runHTTPInvalidToolHeadersClient) + + authScenarios := []string{ + "auth/2025-03-26-oauth-metadata-backcompat", + "auth/2025-03-26-oauth-endpoint-fallback", + "auth/authorization-server-migration", + "auth/basic-cimd", + "auth/iss-normalized", + "auth/iss-not-advertised", + "auth/iss-supported", + "auth/iss-supported-missing", + "auth/iss-unexpected", + "auth/iss-wrong-issuer", + "auth/metadata-default", + "auth/metadata-issuer-mismatch", + "auth/metadata-var1", + "auth/metadata-var2", + "auth/metadata-var3", + "auth/offline-access-not-supported", + "auth/offline-access-scope", + "auth/pre-registration", + "auth/resource-mismatch", + "auth/scope-from-scopes-supported", + "auth/scope-from-www-authenticate", + "auth/scope-omitted-when-undefined", + "auth/scope-retry-limit", + "auth/scope-step-up", + "auth/token-endpoint-auth-basic", + "auth/token-endpoint-auth-none", + "auth/token-endpoint-auth-post", + } + for _, scenario := range authScenarios { + registerScenario(scenario, runAuthClient) + } +} + +// ============================================================================ +// Basic scenarios +// ============================================================================ + +func runBasicClient(ctx context.Context, serverURL string, _ map[string]any) error { + session, err := connectToServer(ctx, serverURL) + if err != nil { + return err + } + defer session.Close() + + _, err = session.ListTools(ctx, nil) + if err != nil { + return fmt.Errorf("session.ListTools(): %v", err) + } + + return nil +} + +func runToolsCallClient(ctx context.Context, serverURL string, _ map[string]any) error { + session, err := connectToServer(ctx, serverURL) + if err != nil { + return err + } + defer session.Close() + + tools, err := session.ListTools(ctx, nil) + if err != nil { + return fmt.Errorf("session.ListTools(): %v", err) + } + + idx := slices.IndexFunc(tools.Tools, func(t *mcp.Tool) bool { + return t.Name == "add_numbers" + }) + if idx == -1 { + return fmt.Errorf("tool %q not found", "add_numbers") + } + + _, err = session.CallTool(ctx, &mcp.CallToolParams{ + Name: "add_numbers", + Arguments: map[string]any{"a": 5, "b": 3}, + }) + if err != nil { + return fmt.Errorf("session.CallTool('add_numbers'): %v", err) + } + + return nil +} + +// ============================================================================ +// Elicitation scenarios +// ============================================================================ + +func runElicitationDefaultsClient(ctx context.Context, serverURL string, _ map[string]any) error { + elicitationHandler := func(ctx context.Context, req *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + return &mcp.ElicitResult{ + Action: "accept", + Content: map[string]any{}, + }, nil + } + + session, err := connectToServer(ctx, serverURL, withClientOptions(&mcp.ClientOptions{ + ElicitationHandler: elicitationHandler, + })) + if err != nil { + return err + } + defer session.Close() + + tools, err := session.ListTools(ctx, nil) + if err != nil { + return fmt.Errorf("session.ListTools(): %v", err) + } + + var testToolName = "test_client_elicitation_defaults" + idx := slices.IndexFunc(tools.Tools, func(t *mcp.Tool) bool { + return t.Name == testToolName + }) + if idx == -1 { + return fmt.Errorf("tool %q not found", testToolName) + } + + _, err = session.CallTool(ctx, &mcp.CallToolParams{ + Name: testToolName, + Arguments: map[string]any{}, + }) + if err != nil { + return fmt.Errorf("session.CallTool(%q): %v", testToolName, err) + } + + return nil +} + +// ============================================================================ +// SSE retry scenario +// ============================================================================ + +func runSSERetryClient(ctx context.Context, serverURL string, _ map[string]any) error { + session, err := connectToServer(ctx, serverURL) + if err != nil { + return err + } + defer session.Close() + log.Printf("Connected to server %q", serverURL) + + tools, err := session.ListTools(ctx, nil) + if err != nil { + return fmt.Errorf("session.ListTools(): %v", err) + } + + var testToolName = "test_reconnection" + idx := slices.IndexFunc(tools.Tools, func(t *mcp.Tool) bool { + return t.Name == testToolName + }) + if idx == -1 { + return fmt.Errorf("tool %q not found", testToolName) + } + + _, err = session.CallTool(ctx, &mcp.CallToolParams{ + Name: testToolName, + Arguments: map[string]any{}, + }) + if err != nil { + return fmt.Errorf("session.CallTool(%q): %v", testToolName, err) + } + + return nil +} + +// ============================================================================ +// Auth scenarios +// ============================================================================ + +func fetchAuthorizationCodeAndState(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + client := &http.Client{ + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } + req, err := http.NewRequestWithContext(ctx, "GET", args.URL, nil) + if err != nil { + return nil, err + } + + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + // In conformance tests the authorization server immediately redirects + // to the callback URL with the authorization code and state. + locURL, err := url.Parse(resp.Header.Get("Location")) + if err != nil { + return nil, fmt.Errorf("parse location: %v", err) + } + + return &auth.AuthorizationResult{ + Code: locURL.Query().Get("code"), + State: locURL.Query().Get("state"), + Iss: locURL.Query().Get("iss"), + }, nil +} + +func runAuthClient(ctx context.Context, serverURL string, configCtx map[string]any) error { + authConfig := &auth.AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost:3000/callback", + AuthorizationCodeFetcher: fetchAuthorizationCodeAndState, + // Try client ID metadata document based registration. + ClientIDMetadataDocumentConfig: &auth.ClientIDMetadataDocumentConfig{ + URL: "https://conformance-test.local/client-metadata.json", + }, + // Try dynamic client registration. + DynamicClientRegistrationConfig: &auth.DynamicClientRegistrationConfig{ + Metadata: &oauthex.ClientRegistrationMetadata{ + RedirectURIs: []string{"http://localhost:3000/callback"}, + }, + }, + } + // Try pre-registered client information if provided in the context. + if clientID, ok := configCtx["client_id"].(string); ok { + if clientSecret, ok := configCtx["client_secret"].(string); ok { + authConfig.PreregisteredClient = &oauthex.ClientCredentials{ + ClientID: clientID, + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: clientSecret, + }, + } + } + } + + authHandler, err := auth.NewAuthorizationCodeHandler(authConfig) + if err != nil { + return fmt.Errorf("failed to create auth handler: %w", err) + } + + session, err := connectToServer(ctx, serverURL, withOAuthHandler(authHandler)) + if err != nil { + return err + } + defer session.Close() + + if _, err := session.ListTools(ctx, nil); err != nil { + return fmt.Errorf("session.ListTools(): %v", err) + } + + if _, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: "test-tool", + Arguments: map[string]any{}, + }); err != nil { + return fmt.Errorf("session.CallTool('test-tool'): %v", err) + } + + return nil +} + +// ============================================================================ +// request-metadata scenario (SEP-2575) +// ============================================================================ + +// runRequestMetadataClient exercises the SEP-2575 wire-level negotiation: +// every request must carry the MCP-Protocol-Version header and the per-request +// _meta envelope, and the client must retry with a supported version when its +// first choice is rejected with -32022. The Go SDK's Connect() drives +// server/discover unconditionally for 2026-07-28, which is exactly that +// mechanism. +func runRequestMetadataClient(ctx context.Context, serverURL string, _ map[string]any) error { + session, err := connectToServer(ctx, serverURL) + if err != nil { + return err + } + defer session.Close() + return nil +} + +// ============================================================================ +// json-schema-ref-no-deref scenario (SEP-2106) +// ============================================================================ + +// runJSONSchemaRefNoDerefClient asserts that listTools does not dereference +// network $ref URLs in tool schemas. The Go SDK never fetches external refs; +// a plain connect → listTools is sufficient — if the SDK ever regressed and +// tried to GET the canary URL the conformance referee would record the +// failure. +func runJSONSchemaRefNoDerefClient(ctx context.Context, serverURL string, _ map[string]any) error { + session, err := connectToServer(ctx, serverURL) + if err != nil { + return err + } + defer session.Close() + + _, _ = session.ListTools(ctx, nil) + return nil +} + +// ============================================================================ +// json-schema-2020-12-preservation scenario (SEP-1613, SEP-2106) +// ============================================================================ + +// runJSONSchemaPreservationClient echoes the focal tool's inputSchema back to +// the referee verbatim, through the permissive json_schema_echo tool, so that +// it can observe which 2020-12 keywords survived the client's parse of +// tools/list. The Go SDK surfaces a server's schema as the raw JSON object it +// received, so nothing is stripped. +func runJSONSchemaPreservationClient(ctx context.Context, serverURL string, _ map[string]any) error { + session, err := connectToServer(ctx, serverURL) + if err != nil { + return err + } + defer session.Close() + + tools, err := session.ListTools(ctx, nil) + if err != nil { + return fmt.Errorf("session.ListTools(): %v", err) + } + + const focalTool = "json_schema_2020_12_tool" + idx := slices.IndexFunc(tools.Tools, func(t *mcp.Tool) bool { + return t.Name == focalTool + }) + if idx == -1 { + return fmt.Errorf("tool %q not found", focalTool) + } + + const echoTool = "json_schema_echo" + if _, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: echoTool, + Arguments: map[string]any{"schema": tools.Tools[idx].InputSchema}, + }); err != nil { + return fmt.Errorf("session.CallTool(%q): %v", echoTool, err) + } + + return nil +} + +// ============================================================================ +// SEP-2322 multi-round-trip client scenario +// ============================================================================ + +// runMrtrClient drives the SEP-2322 client-request-state scenario. The +// SDK's built-in client-side MRTR middleware transparently fulfills embedded +// elicitation requests via ElicitationHandler and echoes requestState. +func runMrtrClient(ctx context.Context, serverURL string, _ map[string]any) error { + elicitationHandler := func(context.Context, *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + return &mcp.ElicitResult{ + Action: "accept", + Content: map[string]any{"confirmed": true}, + }, nil + } + session, err := connectToServer(ctx, serverURL, + withClientOptions(&mcp.ClientOptions{ + Capabilities: &mcp.ClientCapabilities{Elicitation: &mcp.ElicitationCapabilities{}}, + ElicitationHandler: elicitationHandler, + }), + ) + if err != nil { + return err + } + defer session.Close() + + toolNames := []string{ + "test_mrtr_echo_state", + "test_mrtr_no_state", + "test_mrtr_unrelated", + "test_mrtr_no_result_type", + } + for _, name := range toolNames { + if _, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: name, + Arguments: map[string]any{}, + }); err != nil { + log.Printf("CallTool(%q) rejected locally (expected for some cases): %v", name, err) + } + } + return nil +} + +// ============================================================================ +// SEP-2243 HTTP header scenarios +// ============================================================================ + +// runHTTPStandardHeadersClient (SEP-2243 standard headers) drives tools, +// resources and prompts once each so the referee sees Mcp-Method / Mcp-Name +// headers on every verb the SDK supports. +func runHTTPStandardHeadersClient(ctx context.Context, serverURL string, _ map[string]any) error { + session, err := connectToServer(ctx, serverURL) + if err != nil { + return err + } + defer session.Close() + + if tools, err := session.ListTools(ctx, nil); err == nil && len(tools.Tools) > 0 { + _, _ = session.CallTool(ctx, &mcp.CallToolParams{ + Name: tools.Tools[0].Name, + Arguments: map[string]any{}, + }) + } + if res, err := session.ListResources(ctx, nil); err == nil && len(res.Resources) > 0 { + _, _ = session.ReadResource(ctx, &mcp.ReadResourceParams{URI: res.Resources[0].URI}) + } + if prompts, err := session.ListPrompts(ctx, nil); err == nil && len(prompts.Prompts) > 0 { + _, _ = session.GetPrompt(ctx, &mcp.GetPromptParams{ + Name: prompts.Prompts[0].Name, + Arguments: map[string]string{}, + }) + } + return nil +} + +// runHTTPCustomHeadersClient (SEP-2243 custom headers) lists tools (so the +// SDK caches inputSchema and its x-mcp-header annotations), then makes the +// runner-supplied tool calls so the referee validates the Mcp-Param-* +// headers. +func runHTTPCustomHeadersClient(ctx context.Context, serverURL string, configCtx map[string]any) error { + session, err := connectToServer(ctx, serverURL) + if err != nil { + return err + } + defer session.Close() + + if _, err := session.ListTools(ctx, nil); err != nil { + return fmt.Errorf("session.ListTools(): %v", err) + } + for _, call := range readToolCallsContext(configCtx) { + if _, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: call.Name, + Arguments: call.Arguments, + }); err != nil { + log.Printf("CallTool(%q): %v", call.Name, err) + } + } + return nil +} + +// runHTTPInvalidToolHeadersClient (SEP-2243 invalid-tool filtering) lists +// tools — a correct SDK leaves only the valid ones — then calls each survivor +// so the referee records SUCCESS for every excluded tool never called. +func runHTTPInvalidToolHeadersClient(ctx context.Context, serverURL string, _ map[string]any) error { + session, err := connectToServer(ctx, serverURL) + if err != nil { + return err + } + defer session.Close() + + tools, err := session.ListTools(ctx, nil) + if err != nil { + return fmt.Errorf("session.ListTools(): %v", err) + } + for _, tool := range tools.Tools { + _, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: tool.Name, + Arguments: map[string]any{"region": "us-west1"}, + }) + if err != nil { + log.Printf("call %q rejected: %v", tool.Name, err) + } + } + return nil +} + +type toolCall struct { + Name string + Arguments map[string]any +} + +// readToolCallsContext parses the `toolCalls` array from MCP_CONFORMANCE_CONTEXT. +// Format: {"toolCalls": [{"name":"...", "arguments":{...}}, ...]}. +func readToolCallsContext(configCtx map[string]any) []toolCall { + raw, ok := configCtx["toolCalls"].([]any) + if !ok { + return nil + } + out := make([]toolCall, 0, len(raw)) + for _, item := range raw { + obj, ok := item.(map[string]any) + if !ok { + continue + } + name, _ := obj["name"].(string) + args, _ := obj["arguments"].(map[string]any) + if args == nil { + args = map[string]any{} + } + out = append(out, toolCall{Name: name, Arguments: args}) + } + return out +} + +// ============================================================================ +// Main entry point +// ============================================================================ + +func main() { + if len(os.Args) != 2 { + printUsageAndExit("Usage: %s ", os.Args[0]) + } + + serverURL := os.Args[1] + scenarioName := os.Getenv("MCP_CONFORMANCE_SCENARIO") + configCtx := getConformanceContext() + + if scenarioName == "" { + printUsageAndExit("MCP_CONFORMANCE_SCENARIO not set") + } + + handler, ok := registry[scenarioName] + if !ok { + printUsageAndExit("Unknown scenario: %q", scenarioName) + } + + ctx := context.Background() + if err := handler(ctx, serverURL, configCtx); err != nil { + log.Fatalf("Scenario %q failed: %v", scenarioName, err) + } +} + +func getConformanceContext() map[string]any { + ctxStr := os.Getenv("MCP_CONFORMANCE_CONTEXT") + if ctxStr == "" { + return nil + } + var ctx map[string]any + _ = json.Unmarshal([]byte(ctxStr), &ctx) + return ctx +} + +func printUsageAndExit(format string, args ...any) { + var scenarios []string + for name := range registry { + scenarios = append(scenarios, name) + } + sort.Strings(scenarios) + + msg := fmt.Sprintf(format, args...) + log.Fatalf("%s\nAvailable scenarios:\n - %s", msg, strings.Join(scenarios, "\n - ")) +} + +type connectConfig struct { + clientOptions *mcp.ClientOptions + oauthHandler auth.OAuthHandler +} + +type connectOption func(*connectConfig) + +func withClientOptions(opts *mcp.ClientOptions) connectOption { + return func(c *connectConfig) { + c.clientOptions = opts + } +} + +func withOAuthHandler(handler auth.OAuthHandler) connectOption { + return func(c *connectConfig) { + c.oauthHandler = handler + } +} + +// connectToServer connects to the MCP server and returns a client session. +// The caller is responsible for closing the session. +func connectToServer(ctx context.Context, serverURL string, opts ...connectOption) (*mcp.ClientSession, error) { + config := &connectConfig{} + for _, opt := range opts { + opt(config) + } + + client := mcp.NewClient(&mcp.Implementation{ + Name: "test-client", + Version: "1.0.0", + }, config.clientOptions) + + transport := &mcp.StreamableClientTransport{ + Endpoint: serverURL, + OAuthHandler: config.oauthHandler, + } + + session, err := client.Connect(ctx, transport, nil) + if err != nil { + return nil, fmt.Errorf("client.Connect(): %w", err) + } + + return session, nil +} diff --git a/conformance/everything-server/main.go b/conformance/everything-server/main.go new file mode 100644 index 0000000..c532d43 --- /dev/null +++ b/conformance/everything-server/main.go @@ -0,0 +1,1316 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +// The conformance server implements features required for MCP conformance testing. +// It mirrors the functionality of the TypeScript conformance server at +// https://github.com/modelcontextprotocol/conformance/blob/main/examples/servers/typescript/everything-server.ts +// +//lint:file-ignore SA1019 conformance server exercises deprecated SEP-2577 APIs (roots, sampling, logging). +package main + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "flag" + "fmt" + "log" + "net/http" + "os" + "strings" + "time" + + "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/yosida95/uritemplate/v3" +) + +var ( + httpAddr = flag.String("http", "", "if set, use streamable HTTP at this address, instead of stdin/stdout") + stateless = flag.Bool("stateless", true, "use stateless streamable HTTP mode") +) + +const watchedResourceURI = "test://watched-resource" + +func main() { + flag.Parse() + + opts := &mcp.ServerOptions{ + CompletionHandler: completionHandler, + SubscribeHandler: subscribeHandler, + UnsubscribeHandler: unsubscribeHandler, + } + + server := mcp.NewServer(&mcp.Implementation{ + Name: "mcp-conformance-test-server", + Version: "1.0.0", + }, opts) + + // Register server features. + registerTools(server) + registerResources(server) + registerPrompts(server) + + // Start the watched resource auto-update goroutine. + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go watchedResourceUpdater(ctx, server) + + // Serve over stdio, or streamable HTTP if -http is set. + if *httpAddr != "" { + handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { + return server + }, &mcp.StreamableHTTPOptions{Stateless: *stateless}) + log.Printf("Conformance server listening at %s", *httpAddr) + log.Fatal(http.ListenAndServe(*httpAddr, handler)) + } else { + t := &mcp.StdioTransport{} + if err := server.Run(ctx, t); err != nil { + log.Printf("Server failed: %v", err) + os.Exit(1) + } + } +} + +// watchedResourceUpdater sends resource update notifications every 3 seconds. +func watchedResourceUpdater(ctx context.Context, server *mcp.Server) { + ticker := time.NewTicker(3 * time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + server.ResourceUpdated(ctx, &mcp.ResourceUpdatedNotificationParams{ + URI: watchedResourceURI, + }) + } + } +} + +// ============================================================================= +// Tools +// ============================================================================= + +func registerTools(server *mcp.Server) { + mcp.AddTool(server, &mcp.Tool{ + Name: "test_simple_text", + Description: "Tests simple text content response", + }, testSimpleTextHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_image_content", + Description: "Tests image content response", + }, testImageContentHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_audio_content", + Description: "Tests audio content response", + }, testAudioContentHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_embedded_resource", + Description: "Tests embedded resource content response", + }, testEmbeddedResourceHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_multiple_content_types", + Description: "Tests response with multiple content types (text, image, resource)", + }, testMultipleContentTypesHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_tool_with_logging", + Description: "Tests tool that emits log messages during execution", + }, testToolWithLoggingHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_tool_with_progress", + Description: "Tests tool that reports progress notifications", + }, testToolWithProgressHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_error_handling", + Description: "Tests error response handling", + }, testErrorHandlingHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_sampling", + Description: "Tests server-initiated sampling (LLM completion request)", + }, testSamplingHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_elicitation", + Description: "Tests server-initiated elicitation (user input request)", + }, testElicitationHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_elicitation_sep1034_defaults", + Description: "Tests elicitation with default values per SEP-1034", + }, testElicitationDefaultsHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_elicitation_sep1330_enums", + Description: "Tests elicitation with enum schema improvements per SEP-1330", + }, testElicitationEnumsHandler) + + // SEP-1613 / SEP-2106: JSON Schema 2020-12 conformance test tool. + // The scenario verifies that $schema/$defs/additionalProperties (SEP-1613) + // and the broader 2020-12 vocabulary — $anchor, allOf/anyOf, if/then/else — + // (SEP-2106) survive tools/list verbatim. + mcp.AddTool(server, &mcp.Tool{ + Name: "json_schema_2020_12_tool", + Description: "Tool with JSON Schema 2020-12 features for conformance testing (SEP-1613, SEP-2106)", + InputSchema: json.RawMessage(`{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "$defs": { + "address": { + "$anchor": "addressDef", + "type": "object", + "properties": { + "street": { "type": "string" }, + "city": { "type": "string" } + } + } + }, + "properties": { + "name": { "type": "string" }, + "address": { "$ref": "#/$defs/address" }, + "contactMethod": { "type": "string", "enum": ["phone", "email"] }, + "phone": { "type": "string" }, + "email": { "type": "string" } + }, + "allOf": [ + { "anyOf": [{ "required": ["phone"] }, { "required": ["email"] }] } + ], + "if": { + "properties": { "contactMethod": { "const": "phone" } }, + "required": ["contactMethod"] + }, + "then": { "required": ["phone"] }, + "else": { "required": ["email"] }, + "additionalProperties": false + }`), + }, jsonSchema202012Handler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_reconnection", + Description: "Tests SSE stream disconnection and client reconnection (SEP-1699). Server will close the stream mid-call and send the result after client reconnects.", + }, testReconnectionHandler) + + // SEP-2243 x-mcp-header tool — arms the http-custom-header-server-validation + // conformance scenario (which skips when no tool with an x-mcp-header + // annotation is found). + mcp.AddTool(server, &mcp.Tool{ + Name: "test_x_mcp_header", + Description: "Tests SEP-2243 Mcp-Param-* server-side validation", + InputSchema: json.RawMessage(`{ + "type": "object", + "properties": { + "region": { "type": "string", "description": "mirrored into Mcp-Param-Region", "x-mcp-header": "Region" }, + "level": { "type": "integer", "description": "non-mirrored argument" } + } + }`), + }, testXMcpHeaderHandler) + + // SEP-2575 diagnostic tools used by the stateless conformance scenario. + mcp.AddTool(server, &mcp.Tool{ + Name: "test_missing_capability", + Description: "Requires the sampling client capability; used to verify MissingRequiredClientCapabilityError (-32021) (SEP-2575)", + }, testMissingCapabilityHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_streaming_elicitation", + Description: "Streams progress notifications while a call is in flight; used to verify response streams carry only notifications, never independent JSON-RPC requests (SEP-2575)", + }, testStreamingElicitationHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_logging_tool", + Description: "Attempts to emit a log message; the framework must drop it when the client did not set _meta.../logLevel (SEP-2575)", + }, testLoggingToolHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_trigger_tool_change", + Description: "Mutates the tool list to trigger a notifications/tools/list_changed on active subscription streams (SEP-2575)", + }, makeTriggerToolChangeHandler(server)) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_trigger_prompt_change", + Description: "Mutates the prompt list to trigger a notifications/prompts/list_changed on active subscription streams (SEP-2575)", + }, makeTriggerPromptChangeHandler(server)) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_input_required_result_elicitation", + Description: "MRTR (SEP-2322): asks for the caller name via an in-band elicitation request", + }, testInputRequiredElicitationHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_input_required_result_sampling", + Description: "MRTR (SEP-2322): asks for an LLM completion via an in-band sampling request", + }, testInputRequiredSamplingHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_input_required_result_list_roots", + Description: "MRTR (SEP-2322): asks for the client roots via an in-band roots/list request", + }, testInputRequiredListRootsHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_input_required_result_request_state", + Description: "MRTR (SEP-2322): round-trips integrity-protected requestState alongside an elicitation request", + }, testInputRequiredRequestStateHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_input_required_result_multiple_inputs", + Description: "MRTR (SEP-2322): asks for elicitation, sampling and roots input in a single round", + }, testInputRequiredMultipleInputsHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_input_required_result_multi_round", + Description: "MRTR (SEP-2322): two elicitation rounds with evolving requestState before completing", + }, testInputRequiredMultiRoundHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_input_required_result_tampered_state", + Description: "MRTR (SEP-2322): rejects retries whose requestState fails integrity verification", + }, testInputRequiredTamperedStateHandler) + + mcp.AddTool(server, &mcp.Tool{ + Name: "test_input_required_result_capabilities", + Description: "MRTR (SEP-2322): only requests input kinds the declared client capabilities cover", + }, testInputRequiredCapabilitiesHandler) +} + +// Tool handlers + +func testSimpleTextHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "This is a simple text response for testing."}, + }, + }, nil, nil +} + +func testImageContentHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.ImageContent{ + Data: imageData(), + MIMEType: "image/png", + }, + }, + }, nil, nil +} + +func testAudioContentHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.AudioContent{ + Data: audioData(), + MIMEType: "audio/wav", + }, + }, + }, nil, nil +} + +func testEmbeddedResourceHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.EmbeddedResource{ + Resource: &mcp.ResourceContents{ + URI: "test://embedded-resource", + MIMEType: "text/plain", + Text: "This is an embedded resource", + }, + }, + }, + }, nil, nil +} + +func testMultipleContentTypesHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "This is text content"}, + &mcp.ImageContent{ + Data: imageData(), + MIMEType: "image/png", + }, + &mcp.EmbeddedResource{ + Resource: &mcp.ResourceContents{ + URI: "test://embedded-in-multiple", + MIMEType: "text/plain", + Text: "This is an embedded resource", + }, + }, + }, + }, nil, nil +} + +func testToolWithLoggingHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + // Emit three info-level log messages + req.Session.Log(ctx, &mcp.LoggingMessageParams{ + Level: "info", + Data: "Tool execution started", + }) + time.Sleep(50 * time.Millisecond) + req.Session.Log(ctx, &mcp.LoggingMessageParams{ + Level: "info", + Data: "Tool processing data", + }) + time.Sleep(50 * time.Millisecond) + req.Session.Log(ctx, &mcp.LoggingMessageParams{ + Level: "info", + Data: "Tool execution completed", + }) + + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "Tool with logging executed successfully"}, + }, + }, nil, nil +} + +func testToolWithProgressHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + // Get progress token from the request if provided + progressToken := req.Params.GetProgressToken() + + // Send three progress notifications (0%, 50%, 100%) + total := 100.0 + steps := []float64{0, 50, 100} + for _, progress := range steps { + req.Session.NotifyProgress(ctx, &mcp.ProgressNotificationParams{ + ProgressToken: progressToken, + Progress: progress, + Total: total, + Message: fmt.Sprintf("Completed step %.0f of %.0f", progress, total), + }) + time.Sleep(50 * time.Millisecond) + } + + // Return the progress token value as the response (matching TypeScript behavior) + tokenStr := fmt.Sprintf("%v", progressToken) + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: tokenStr}, + }, + }, nil, nil +} + +func testErrorHandlingHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + return nil, nil, errors.New("this tool intentionally returns an error for testing") +} + +type samplingInput struct { + Prompt string `json:"prompt" jsonschema:"The prompt to send to the LLM"` +} + +func testSamplingHandler(ctx context.Context, req *mcp.CallToolRequest, input samplingInput) (*mcp.CallToolResult, any, error) { + // Request LLM completion from the client + result, err := req.Session.CreateMessage(ctx, &mcp.CreateMessageParams{ + Messages: []*mcp.SamplingMessage{ + { + Role: "user", + Content: &mcp.TextContent{ + Text: input.Prompt, + }, + }, + }, + MaxTokens: 100, + }) + if err != nil { + return nil, nil, fmt.Errorf("sampling failed: %w", err) + } + + // Extract the text response from the result + var responseText string + if tc, ok := result.Content.(*mcp.TextContent); ok { + responseText = tc.Text + } else { + responseText = "(non-text response)" + } + + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: fmt.Sprintf("LLM response: %s", responseText)}, + }, + }, nil, nil +} + +type elicitationInput struct { + Message string `json:"message" jsonschema:"The message to show the user"` +} + +func testElicitationHandler(ctx context.Context, req *mcp.CallToolRequest, input elicitationInput) (*mcp.CallToolResult, any, error) { + result, err := req.Session.Elicit(ctx, &mcp.ElicitParams{ + Message: input.Message, + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "username": { + Type: "string", + Description: "Your preferred username", + }, + }, + Required: []string{"username"}, + }, + }) + if err != nil { + return nil, nil, fmt.Errorf("elicitation failed: %w", err) + } + + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: fmt.Sprintf("Elicitation result: action=%s, content=%v", result.Action, result.Content)}, + }, + }, nil, nil +} + +func testElicitationDefaultsHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + result, err := req.Session.Elicit(ctx, &mcp.ElicitParams{ + Message: "Test defaults for primitives", + RequestedSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": map[string]any{ + "type": "string", + "description": "User name", + "default": "John Doe", + }, + "age": map[string]any{ + "type": "integer", + "description": "User age", + "default": 30, + }, + "score": map[string]any{ + "type": "number", + "description": "User score", + "default": 95.5, + }, + "status": map[string]any{ + "type": "string", + "description": "User status", + "enum": []string{"active", "inactive", "pending"}, + "default": "active", + }, + "verified": map[string]any{ + "type": "boolean", + "description": "Verification status", + "default": true, + }, + }, + }, + }) + if err != nil { + return nil, nil, fmt.Errorf("elicitation failed: %w", err) + } + + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: fmt.Sprintf("Elicitation result: action=%s, content=%v", result.Action, result.Content)}, + }, + }, nil, nil +} + +func testElicitationEnumsHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + result, err := req.Session.Elicit(ctx, &mcp.ElicitParams{ + Message: "Test enum schemas", + RequestedSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + // Basic enum without titles + "untitledSingle": map[string]any{ + "type": "string", + "enum": []string{"option1", "option2", "option3"}, + }, + // Enum with titles using oneOf + "titledSingle": map[string]any{ + "type": "string", + "oneOf": []map[string]any{ + {"const": "value1", "title": "First Option"}, + {"const": "value2", "title": "Second Option"}, + {"const": "value3", "title": "Third Option"}, + }, + }, + // Legacy enum with enumNames + "legacyEnum": map[string]any{ + "type": "string", + "enum": []string{"opt1", "opt2", "opt3"}, + "enumNames": []string{"Option One", "Option Two", "Option Three"}, + }, + // Multi-select without titles + "untitledMulti": map[string]any{ + "type": "array", + "minItems": 1, + "maxItems": 3, + "items": map[string]any{ + "type": "string", + "enum": []string{"option1", "option2", "option3"}, + }, + }, + // Multi-select with titles using anyOf + "titledMulti": map[string]any{ + "type": "array", + "minItems": 1, + "maxItems": 3, + "items": map[string]any{ + "type": "string", + "anyOf": []map[string]any{ + {"const": "value1", "title": "First Option"}, + {"const": "value2", "title": "Second Option"}, + {"const": "value3", "title": "Third Option"}, + }, + }, + }, + }, + }, + }) + if err != nil { + return nil, nil, fmt.Errorf("elicitation failed: %w", err) + } + + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: fmt.Sprintf("Elicitation result: action=%s, content=%v", result.Action, result.Content)}, + }, + }, nil, nil +} + +func jsonSchema202012Handler(ctx context.Context, req *mcp.CallToolRequest, input json.RawMessage) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{ + Text: fmt.Sprintf("JSON Schema 2020-12 tool called with: %s", input), + }, + }, + }, nil, nil +} + +type testXMcpHeaderInput struct { + Region string `json:"region,omitempty"` + Level int `json:"level,omitempty"` +} + +func testXMcpHeaderHandler(ctx context.Context, req *mcp.CallToolRequest, input testXMcpHeaderInput) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: fmt.Sprintf("region=%s", input.Region)}, + }, + }, nil, nil +} + +// testMissingCapabilityHandler requires the client to have declared the +// sampling capability. When absent, it returns a MissingRequiredClientCapabilityError +// with data.requiredCapabilities = {"sampling": {}} per SEP-2575. +func testMissingCapabilityHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + caps := req.ClientCapabilities() + if caps != nil && caps.Sampling != nil { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "Client declared the sampling capability; tool executed."}, + }, + }, nil, nil + } + missingCapabilityData := mcp.MissingRequiredClientCapabilityData{ + RequiredCapabilities: &mcp.ClientCapabilities{ + Sampling: &mcp.SamplingCapabilities{}, + }, + } + dataBytes, err := json.Marshal(missingCapabilityData) + if err != nil { + return nil, nil, err + } + return nil, nil, &jsonrpc.Error{ + Code: mcp.CodeMissingRequiredClientCapabilities, + Message: "sampling capability required but not declared by client", + Data: dataBytes, + } +} + +// testStreamingElicitationHandler returns a plain result. Per SEP-2575, the +// response stream must carry no independent top-level JSON-RPC requests; a +// plain response trivially satisfies this. The scenario declares no +// `elicitation` capability, so this handler must not call req.Session.Elicit. +func testStreamingElicitationHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "stream observed: result frames only, no top-level requests"}, + }, + }, nil, nil +} + +// testLoggingToolHandler attempts to emit a log message. The SDK's +// ServerSession.Log gates on the per-request _meta.../logLevel, so this is a +// no-op notification-wise when the client did not opt in. +func testLoggingToolHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + req.Session.Log(ctx, &mcp.LoggingMessageParams{ + Level: "info", + Data: "test_logging_tool executed", + }) + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "Log attempted; framework gates on _meta.../logLevel."}, + }, + }, nil, nil +} + +// makeTriggerToolChangeHandler returns a handler that (re-)registers a +// no-op transient tool. Every call to Server.AddTool dispatches a +// notifications/tools/list_changed to active subscription streams (even when +// it replaces a tool with the same name), which is what the SEP-2575 +// subscription checks assert. +func makeTriggerToolChangeHandler(server *mcp.Server) mcp.ToolHandlerFor[any, any] { + return func(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + mcp.AddTool(server, &mcp.Tool{ + Name: "__transient_tool_for_list_changed", + Description: "Transient tool used to trigger tools/list_changed", + }, func(context.Context, *mcp.CallToolRequest, any) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{}, nil, nil + }) + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "tools_list_changed published"}, + }, + }, nil, nil + } +} + +// makeTriggerPromptChangeHandler is the prompt analogue of +// makeTriggerToolChangeHandler. +func makeTriggerPromptChangeHandler(server *mcp.Server) mcp.ToolHandlerFor[any, any] { + return func(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + server.AddPrompt(&mcp.Prompt{ + Name: "__transient_prompt_for_list_changed", + Description: "Transient prompt used to trigger prompts/list_changed", + }, func(context.Context, *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{}, nil + }) + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "prompts_list_changed published"}, + }, + }, nil, nil + } +} + +// ============================================================================= +// SEP-2322 Multi-Round-Trip helpers +// ============================================================================= + +// requestStateTamperedSuffix is what the tampered-state conformance scenario +// appends to a valid requestState. Detecting this suffix is sufficient to +// pass the check; a production server would use a real integrity mechanism +// (HMAC / AEAD, per SEP-2322 § "Server Requirements"). +const requestStateTamperedSuffix = "-TAMPERED" + +// requestStateValid reports whether the state looks like one this server +// minted (i.e. the harness did not append the tampered marker). +func requestStateValid(state string) bool { + return state != "" && !strings.HasSuffix(state, requestStateTamperedSuffix) +} + +// acceptedContent returns the "content" field of an ElicitResult only when the +// user accepted the elicitation. Returns nil if the response is missing, of the +// wrong type, or the action is not "accept". +func acceptedContent(responses mcp.InputResponseMap, key string) map[string]any { + if responses == nil { + return nil + } + resp, ok := responses[key] + if !ok { + return nil + } + elicit, ok := resp.(*mcp.ElicitResult) + if !ok || elicit == nil || elicit.Action != "accept" { + return nil + } + return elicit.Content +} + +// ============================================================================= +// SEP-2322 Multi-Round-Trip tool handlers +// ============================================================================= + +// nameElicitSchema is the requestedSchema used by tools that elicit a name. +func nameElicitSchema() *jsonschema.Schema { + return &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "name": {Type: "string"}, + }, + Required: []string{"name"}, + } +} + +// testInputRequiredElicitationHandler asks for the caller name via elicitation. +// Reused by the result-type, missing-input-response, ignore-extra-params, and +// validate-input scenarios: anything that does not carry an accepted +// "user_name" response is answered with a fresh InputRequiredResult. +func testInputRequiredElicitationHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + if content := acceptedContent(req.Params.InputResponses, "user_name"); content != nil { + name, _ := content["name"].(string) + if name != "" { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Hello, %s!", name)}}, + }, nil, nil + } + } + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{ + "user_name": &mcp.ElicitParams{ + Message: "What is your name?", + RequestedSchema: nameElicitSchema(), + }, + }, + }, nil, nil +} + +// samplingResponseText extracts the first text content from a sampling +// response. The Go SDK unmarshals responses with a "role" field into +// *CreateMessageWithToolsResult (a superset of the legacy single-content +// CreateMessageResult). +func samplingResponseText(responses mcp.InputResponseMap, key string) (string, bool) { + resp, ok := responses[key].(*mcp.CreateMessageWithToolsResult) + if !ok || resp == nil { + return "", false + } + for _, c := range resp.Content { + if tc, ok := c.(*mcp.TextContent); ok { + return tc.Text, true + } + } + return "(non-text response)", true +} + +// testInputRequiredSamplingHandler asks for an LLM completion via sampling. +func testInputRequiredSamplingHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + if text, ok := samplingResponseText(req.Params.InputResponses, "capital_question"); ok { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Sampling response: %s", text)}}, + }, nil, nil + } + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{ + "capital_question": &mcp.CreateMessageParams{ + Messages: []*mcp.SamplingMessage{ + {Role: "user", Content: &mcp.TextContent{Text: "What is the capital of France?"}}, + }, + MaxTokens: 100, + }, + }, + }, nil, nil +} + +// testInputRequiredListRootsHandler asks for the client roots. +func testInputRequiredListRootsHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + if resp, ok := req.Params.InputResponses["client_roots"].(*mcp.ListRootsResult); ok && resp != nil { + var uris []string + for _, r := range resp.Roots { + uris = append(uris, r.URI) + } + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: fmt.Sprintf("Client exposed %d root(s): %s", len(resp.Roots), strings.Join(uris, ", "))}, + }, + }, nil, nil + } + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{ + "client_roots": &mcp.ListRootsParams{}, + }, + }, nil, nil +} + +// testInputRequiredRequestStateHandler round-trips an opaque requestState +// alongside an elicitation request. +func testInputRequiredRequestStateHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + if content := acceptedContent(req.Params.InputResponses, "confirm"); content != nil { + if requestStateValid(req.Params.RequestState) { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "state-ok: requestState received and confirmation accepted"}}, + }, nil, nil + } + return nil, nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "invalid requestState", + } + } + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{ + "confirm": &mcp.ElicitParams{ + Message: "Please confirm", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{"ok": {Type: "boolean"}}, + Required: []string{"ok"}, + }, + }, + }, + RequestState: "request_state", + }, nil, nil +} + +// testInputRequiredMultipleInputsHandler asks for elicitation, sampling, and +// roots input in a single round. +func testInputRequiredMultipleInputsHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + nameContent := acceptedContent(req.Params.InputResponses, "user_name") + name, _ := nameContent["name"].(string) + + greetingText, hasGreeting := samplingResponseText(req.Params.InputResponses, "greeting") + + var rootsCount int + rootsResp, rootsOK := req.Params.InputResponses["client_roots"].(*mcp.ListRootsResult) + if rootsOK && rootsResp != nil { + rootsCount = len(rootsResp.Roots) + } + + if name != "" && hasGreeting && rootsOK { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: fmt.Sprintf("%s %s — %d root(s) visible", greetingText, name, rootsCount)}, + }, + }, nil, nil + } + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{ + "user_name": &mcp.ElicitParams{ + Message: "What is your name?", + RequestedSchema: nameElicitSchema(), + }, + "greeting": &mcp.CreateMessageParams{ + Messages: []*mcp.SamplingMessage{ + {Role: "user", Content: &mcp.TextContent{Text: "Generate a greeting"}}, + }, + MaxTokens: 50, + }, + "client_roots": &mcp.ListRootsParams{}, + }, + RequestState: "multiple_inputs", + }, nil, nil +} + +// testInputRequiredMultiRoundHandler runs two elicitation rounds with +// evolving requestState before completing. The round number (and the name +// elicited in round 1) lives in the state, since the stateless server keeps +// no per-session memory. +func testInputRequiredMultiRoundHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + state := req.Params.RequestState + // State format: "round=1" or "round=2;name=". + switch { + case state == "": + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{ + "step1": &mcp.ElicitParams{ + Message: "Step 1: What is your name?", + RequestedSchema: nameElicitSchema(), + }, + }, + RequestState: "round=1", + }, nil, nil + case state == "round=1": + name, _ := acceptedContent(req.Params.InputResponses, "step1")["name"].(string) + if name == "" { + name = "unknown" + } + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{ + "step2": &mcp.ElicitParams{ + Message: "Step 2: What is your favorite color?", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{"color": {Type: "string"}}, + Required: []string{"color"}, + }, + }, + }, + RequestState: "round=2;name=" + name, + }, nil, nil + default: + color, _ := acceptedContent(req.Params.InputResponses, "step2")["color"].(string) + if color == "" { + color = "unknown" + } + name := strings.TrimPrefix(state, "round=2;name=") + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: fmt.Sprintf("Multi-round complete: %s likes %s", name, color)}, + }, + }, nil, nil + } +} + +// testInputRequiredTamperedStateHandler rejects retries whose requestState +// looks tampered. The conformance harness tampers by appending "-TAMPERED"; +// production servers should use a real integrity mechanism (HMAC / AEAD). +func testInputRequiredTamperedStateHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + if req.Params.RequestState != "" { + if !requestStateValid(req.Params.RequestState) { + return nil, nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "invalid requestState: tampered", + } + } + if acceptedContent(req.Params.InputResponses, "confirm") != nil { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "state-ok: requestState received and confirmation accepted"}}, + }, nil, nil + } + } + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{ + "confirm": &mcp.ElicitParams{ + Message: "Please confirm", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{"ok": {Type: "boolean"}}, + Required: []string{"ok"}, + }, + }, + }, + RequestState: "tampered_state", + }, nil, nil +} + +// testInputRequiredCapabilitiesHandler only requests input kinds the client +// declared capabilities cover. Reads capabilities from the per-request +// _meta["io.modelcontextprotocol/clientCapabilities"] envelope. +func testInputRequiredCapabilitiesHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) > 0 { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "Capability-aware input requests fulfilled"}}, + }, nil, nil + } + caps := req.ClientCapabilities() + reqs := mcp.InputRequestMap{} + if caps != nil && caps.Elicitation != nil { + reqs["user_name"] = &mcp.ElicitParams{ + Message: "What is your name?", + RequestedSchema: nameElicitSchema(), + } + } + if caps != nil && caps.Sampling != nil { + reqs["greeting"] = &mcp.CreateMessageParams{ + Messages: []*mcp.SamplingMessage{ + {Role: "user", Content: &mcp.TextContent{Text: "Generate a short greeting"}}, + }, + MaxTokens: 50, + } + } + if caps != nil && (caps.RootsV2 != nil || caps.Roots.ListChanged) { + reqs["client_roots"] = &mcp.ListRootsParams{} + } + if len(reqs) == 0 { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "No declared client capability supports an in-band input request"}}, + }, nil, nil + } + return &mcp.CallToolResult{InputRequests: reqs}, nil, nil +} + +// testInputRequiredPromptHandler is the prompts/get analogue: elicits context +// on the first call and returns a GetPromptResult with messages on retry. +func testInputRequiredPromptHandler(ctx context.Context, req *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + if content := acceptedContent(req.Params.InputResponses, "user_context"); content != nil { + contextStr, _ := content["context"].(string) + return &mcp.GetPromptResult{ + Description: "A prompt with elicited context", + Messages: []*mcp.PromptMessage{ + { + Role: "user", + Content: &mcp.TextContent{Text: fmt.Sprintf("Context: %s", contextStr)}, + }, + }, + }, nil + } + return &mcp.GetPromptResult{ + InputRequests: mcp.InputRequestMap{ + "user_context": &mcp.ElicitParams{ + Message: "Please provide context for the prompt", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{"context": {Type: "string"}}, + Required: []string{"context"}, + }, + }, + }, + }, nil +} + +func testReconnectionHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + // Close the SSE stream to trigger client reconnection (SEP-1699) + if req.Extra != nil && req.Extra.CloseSSEStream != nil { + req.Extra.CloseSSEStream(mcp.CloseSSEStreamArgs{RetryAfter: 10 * time.Millisecond}) + } + + // Wait for client to reconnect + time.Sleep(100 * time.Millisecond) + + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{ + Text: "Reconnection test completed successfully. If you received this, the client properly reconnected after stream closure.", + }, + }, + }, nil, nil +} + +// ============================================================================= +// Resources +// ============================================================================= + +func registerResources(server *mcp.Server) { + server.AddResource(&mcp.Resource{ + Name: "static-text", + Description: "A static text resource for testing", + MIMEType: "text/plain", + URI: "test://static-text", + }, staticTextHandler) + + server.AddResource(&mcp.Resource{ + Name: "static-binary", + Description: "A static binary resource (image) for testing", + MIMEType: "image/png", + URI: "test://static-binary", + }, staticBinaryHandler) + + server.AddResourceTemplate(&mcp.ResourceTemplate{ + Name: "template", + Description: "A resource template with parameter substitution", + MIMEType: "application/json", + URITemplate: "test://template/{id}/data", + }, templateResourceHandler) + + server.AddResource(&mcp.Resource{ + Name: "watched-resource", + Description: "A resource that auto-updates every 3 seconds", + MIMEType: "text/plain", + URI: watchedResourceURI, + }, watchedResourceHandler) +} + +func staticTextHandler(ctx context.Context, req *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{ + Contents: []*mcp.ResourceContents{ + { + URI: req.Params.URI, + MIMEType: "text/plain", + Text: "This is the content of the static text resource.", + }, + }, + }, nil +} + +func staticBinaryHandler(ctx context.Context, req *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{ + Contents: []*mcp.ResourceContents{ + { + URI: req.Params.URI, + MIMEType: "image/png", + Blob: imageData(), + }, + }, + }, nil +} + +// templatePattern is the compiled URI template for the template resource. +var templatePattern = uritemplate.MustNew("test://template/{id}/data") + +func templateResourceHandler(ctx context.Context, req *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + // Extract the ID from the URI using the template pattern + uri := req.Params.URI + match := templatePattern.Regexp().FindStringSubmatch(uri) + id := "" + if len(match) > 1 { + id = match[1] + } + + jsonContent := fmt.Sprintf(`{"id": "%s", "templateTest": true, "data": "Data for ID: %s"}`, id, id) + return &mcp.ReadResourceResult{ + Contents: []*mcp.ResourceContents{ + { + URI: uri, + MIMEType: "application/json", + Text: jsonContent, + }, + }, + }, nil +} + +func watchedResourceHandler(ctx context.Context, req *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{ + Contents: []*mcp.ResourceContents{ + { + URI: req.Params.URI, + MIMEType: "text/plain", + Text: "Watched resource content", + }, + }, + }, nil +} + +// ============================================================================= +// Prompts +// ============================================================================= + +func registerPrompts(server *mcp.Server) { + server.AddPrompt(&mcp.Prompt{ + Name: "test_simple_prompt", + Title: "Simple Test Prompt", + Description: "A simple prompt without arguments", + }, simplePromptHandler) + + server.AddPrompt(&mcp.Prompt{ + Name: "test_prompt_with_arguments", + Title: "Prompt With Arguments", + Description: "A prompt with required arguments", + Arguments: []*mcp.PromptArgument{ + {Name: "arg1", Description: "First test argument", Required: true}, + {Name: "arg2", Description: "Second test argument", Required: true}, + }, + }, promptWithArgumentsHandler) + + server.AddPrompt(&mcp.Prompt{ + Name: "test_prompt_with_embedded_resource", + Title: "Prompt With Embedded Resource", + Description: "A prompt that includes an embedded resource", + Arguments: []*mcp.PromptArgument{ + {Name: "resourceUri", Description: "URI of the resource to embed", Required: true}, + }, + }, promptWithEmbeddedResourceHandler) + + server.AddPrompt(&mcp.Prompt{ + Name: "test_prompt_with_image", + Title: "Prompt With Image", + Description: "A prompt that includes image content", + }, promptWithImageHandler) + + server.AddPrompt(&mcp.Prompt{ + Name: "test_input_required_result_prompt", + Description: "MRTR (SEP-2322): asks for context via an in-band elicitation request", + }, testInputRequiredPromptHandler) +} + +func simplePromptHandler(ctx context.Context, req *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{ + Description: "A simple test prompt", + Messages: []*mcp.PromptMessage{ + { + Role: "user", + Content: &mcp.TextContent{Text: "This is a simple prompt for testing."}, + }, + }, + }, nil +} + +func promptWithArgumentsHandler(ctx context.Context, req *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + arg1 := req.Params.Arguments["arg1"] + arg2 := req.Params.Arguments["arg2"] + + return &mcp.GetPromptResult{ + Description: "A prompt with arguments", + Messages: []*mcp.PromptMessage{ + { + Role: "user", + Content: &mcp.TextContent{ + Text: fmt.Sprintf("Prompt with arguments: arg1='%s', arg2='%s'", arg1, arg2), + }, + }, + }, + }, nil +} + +func promptWithEmbeddedResourceHandler(ctx context.Context, req *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + resourceUri := req.Params.Arguments["resourceUri"] + + return &mcp.GetPromptResult{ + Description: "A prompt with an embedded resource", + Messages: []*mcp.PromptMessage{ + { + Role: "user", + Content: &mcp.EmbeddedResource{ + Resource: &mcp.ResourceContents{ + URI: resourceUri, + MIMEType: "text/plain", + Text: "Embedded resource content for testing.", + }, + }, + }, + { + Role: "user", + Content: &mcp.TextContent{Text: "Please process the embedded resource above."}, + }, + }, + }, nil +} + +func promptWithImageHandler(ctx context.Context, req *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{ + Description: "A prompt with an image", + Messages: []*mcp.PromptMessage{ + { + Role: "user", + Content: &mcp.ImageContent{ + Data: imageData(), + MIMEType: "image/png", + }, + }, + { + Role: "user", + Content: &mcp.TextContent{Text: "Please analyze the image above."}, + }, + }, + }, nil +} + +// ============================================================================= +// Server handlers +// ============================================================================= + +func completionHandler(ctx context.Context, req *mcp.CompleteRequest) (*mcp.CompleteResult, error) { + // Return empty completion - just acknowledging the capability + return &mcp.CompleteResult{ + Completion: mcp.CompletionResultDetails{ + Values: []string{}, + Total: 0, + }, + }, nil +} + +func subscribeHandler(ctx context.Context, req *mcp.SubscribeRequest) error { + // The SDK handles subscription tracking internally via Server.ResourceUpdated() + return nil +} + +func unsubscribeHandler(ctx context.Context, req *mcp.UnsubscribeRequest) error { + // The SDK handles subscription tracking internally + return nil +} + +// ============================================================================= +// Helper functions +// ============================================================================= + +// Base64-encoded minimal test files, copied from the typescript conformance example. +const ( + // Minimal 1x1 red PNG image + testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" + // Minimal WAV audio file (silence) + testAudioBase64 = "UklGRiYAAABXQVZFZm10IBAAAAABAAEAQB8AAAB9AAACABAAZGF0YQIAAAA=" +) + +func imageData() []byte { + data, err := base64.StdEncoding.DecodeString(testImageBase64) + if err != nil { + panic("invalid testImageBase64: " + err.Error()) + } + return data +} + +func audioData() []byte { + data, err := base64.StdEncoding.DecodeString(testAudioBase64) + if err != nil { + panic("invalid testAudioBase64: " + err.Error()) + } + return data +} diff --git a/copyright_test.go b/copyright_test.go index 287e594..bda479f 100644 --- a/copyright_test.go +++ b/copyright_test.go @@ -17,8 +17,7 @@ import ( func TestCopyrightHeaders(t *testing.T) { var re = regexp.MustCompile(`Copyright \d{4} The Go MCP SDK Authors. All rights reserved. -Use of this source code is governed by an MIT-style -license that can be found in the LICENSE file.`) +Use of this source code is governed by (the license\n|an MIT-style\nlicense )that can be found in the LICENSE file.`) err := filepath.WalkDir(".", func(path string, d fs.DirEntry, err error) error { if err != nil { diff --git a/design/mrtr.md b/design/mrtr.md new file mode 100644 index 0000000..16321a4 --- /dev/null +++ b/design/mrtr.md @@ -0,0 +1,264 @@ +## Context + +A proposal for implementing Multi Round-Trip Requests +(MRTR) as defined in [SEP-2322](https://github.com/CaitieM20/modelcontextprotocol/blob/de6d76fba3078eda957dadb3cec51ca8ab851b5c/seps/2322-MRTR.md). + +In the new protocol version servers can't initiate requests to clients, but when a server requires additional input for completing `tools/call`, `prompts/get`, or `resources/read` it can return an incomplete result along with a set of `inputRequests`. The client fulfills them locally and retries the same call with `inputResponses` attached. + +## Goals + +**Must have:** + +* Backward compatibility. +* Correct representation on the wire. + +**Nice to have:** + +* Minimal changes to the exported API surface. +* Hard for server implementers to construct an invalid payload. +* Simple input request handling for clients. +* Protocol-version-independent code. +* Consistency with the rest of the SDK. + +## Proposal + +`ServerSession` methods return an error for new-version protocol connections. + +`InputRequest`/`InputResponse` is introduced as a sealed-interface: +```go +// Implemented by *ElicitParams, *CreateMessageParams, *ListRootsParams +type InputRequest interface{ isInputRequest() } + +type InputRequestMap map[string]InputRequest +// MarshalJSON encodes as map[string]struct{ Method string; Params InputRequest } +func (m InputRequestMap) MarshalJSON() ([]byte, error) { ... } +// UnmarshalJSON decodes from map[string]struct{ Method string; Params InputRequest } +func (m *InputRequestMap) UnmarshalJSON(data []byte) error { ... } + +// Implemented by *ElicitResult, *CreateMessageResult, *ListRootsResult. +type InputResponse interface{ isInputResult() } + +type InputResponseMap map[string]InputResponse +// MarshalJSON encodes as map[string]struct{ Method string; Result InputResponse } +func (m InputResponseMap) MarshalJSON() ([]byte, error) { ... } +// UnmarshalJSON decodes from map[string]struct{ Method string; Result InputResponse } +func (m *InputResponseMap) UnmarshalJSON(data []byte) error { ... } +``` + +All affected methods' `*Params` are extended with `InputResponseMap` and `RequestState` fields: +```go +type CallToolParams struct { + ... + InputResponses InputResponseMap `json:"inputResponses,omitempty"` + RequestState string `json:"requestState,omitempty"` +} +// Same for GetPromptParams, ReadResourceParams +``` + +`InputRequests` and `RequestState` fields are added directly to `CallToolResult`, `GetPromptResult`, and `ReadResourceResult` as exported. +Result type discriminator (completed, input_required) is unexported so that SDK users don't need to set it to the correct constant in addition to setting either `Content` or `InputRequests`. Handler execution result is validated and augmented before marshaling: +```go +type CallToolResult struct { + ... + InputRequests InputRequestMap `json:"inputRequests,omitempty"` + RequestState string `json:"requestState,omitempty"` + resultType string // set by the SDK and used in MarshalJSON() +} +// Same for GetPromptResult, ReadResourceResult. +``` +Alternatively, the field could only exist on `wire struct`, but this would make us return `complete` to older clients or empty string to newer clients, because there's no access to negotiated protocol version in `MarshalJSON`. + +Servers request additional input by constructing a correct struct literal: +```go +mcp.AddTool(s, tool, func(ctx context.Context, req *mcp.CallToolRequest, in MyIn) (*mcp.CallToolResult, MyOut, error) { + if !hasConfirmation(in) { + return &mcp.CallToolResult{ + InputRequests: InputRequestMap{"confirm": &mcp.ElicitParams{Message: "Sure?"}}, + RequestState: "state-token", + }, zero, nil + } + return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "done"}}}, myOut, nil +}) +``` +The SDK validates at runtime that a handler does not return both content and `InputRequests` — doing so logs a warning and returns a `CodeInternalError` JSON-RPC error. + +An unexported receiving middleware is installed on the server for backward compatibility with older clients. When a handler returns `InputRequests` and the connected client uses a protocol version that does not support MRTR, the middleware fulfills the requests by calling `ServerSession.Elicit`/`CreateMessage`/`ListRoots` on the client directly and reinvokes the handler once with the collected `InputResponses`. If any of these calls fail, the entire request fails. Input requests are fulfilled concurrently. This lets server developers write protocol-version-independent code. + +An unexported sending middleware is installed on the client, which similarly to `urlElicitationMiddleware` will automatically invoke handlers for the corresponding methods on incomplete results and retry the original request. Clients have an option to disable it and write a retry loop manually using `NeedsInput()`: +```go +type MultiRoundTripOptions struct { + Disabled bool +} + +client := mcp.NewClient(impl, &mcp.ClientOptions{MultiRoundTrip: &mcp.MultiRoundTripOptions{Disabled: true}}) +result, err := client.CallTool(ctx, &mcp.CallToolParams{Name: "my-tool"}) +if result.NeedsInput() { ... } +``` + +`NeedsInput()` checks the unexported `resultType` field rather than `InputRequests`, correctly handling the load-shedding case where the server returns `input_required` with an empty map. + +**Pros** + +This is arguably the simplest and the most transparent approach which is also closest to the spec. +What gets explicitly set on the server can be observed on the wire and on the client. +The opt-out client middleware follows the principle of the least surprise for app developers. If client method handlers were provided they will continue to be invoked regardless of the protocol version in use. The `Disabled` option lets "power-users" build any custom handling logic. +The server middleware makes handler code protocol-version-independent — the same handler works for both old and new clients. + +**Cons** + +The biggest downside of the proposal is that server developers can construct incorrect responses (both content and input requests) and this will only be validated at runtime. + +## Alternatives considered + +### Unexported fields + +MRTR fields can be unexported, accessible only through getters, constructible only through constructor functions, and handled explicitly in custom `(Unm|M)arshalJSON`. This will make it impossible for developers to construct incorrect responses and for clients to perform an erroneous `len(result.InputRequests) > 0` check in the load-shedding case. +```go +type CallToolResult struct { + ... + inputRequests InputRequestMap + requestState string + resultType string +} + +func (r *CallToolResult) InputRequests() (InputRequestMap, bool) { ... } + +// InputRequiredResult struct exists for backward-compatibility in case of new fields being needed for input request results. +type InputRequiredResult struct { + InputRequests InputRequestMap + RequestState string +} + +// RequireInput constructs a tool call, prompt or resource result with input requests set. +// mrtrResult provides methods for setting private fields on these types. +func RequireInput[T any, TP interface { *T; mrtrResult }](r InputRequiredResult) TP { ... } +``` + +On the server: +```go +mcp.AddTool(s, tool, func(ctx context.Context, req *mcp.CallToolRequest, in MyIn) (*mcp.CallToolResult, MyOut, error) { + if !hasConfirmation(in) { + return mcp.RequireInput[mcp.CallToolResult](mcp.InputRequiredResult{ + InputRequests: mcp.InputRequestMap{"confirm": &mcp.ElicitParams{Message: "Deploy to production?"}}, + RequestState: "deployment-123", + }), nil, nil + } + return &mcp.CallToolResult{ Content: []mcp.Content{&mcp.TextContent{Text: "done"}}}, myOut, nil +}) +``` + +On the client: +```go +result, err := client.CallTool(ctx, &mcp.CallToolParams{Name: "my-tool"}) +if requests, ok := result.InputRequests(); ok { ... } +``` + +The biggest downside of this approach is the obscure data model with hidden fields. An incomplete `mcp.CallToolResult` looks like an uninitialized struct until `InputRequests` method result is examined. +In addition to this, the verbose `RequireInput` syntax (no auto type inference from assignment target) does not look idiomatic and fits poorly into the existing SDK APIs. + +--- + +### `InputRequiredError` type + +We could explore a different data channel - `error` return value. This would give us the natural "happy path is when all inputs are provided" flow on the server side, and good result interpretability on the client side (impossible to confuse with a successful response). +The new error could be converted to the correct wire representation at the marshaling stage. +```go +type InputRequiredError struct { + InputRequests InputRequestMap + RequestState string +} + +func (e *InputRequiredError) Error() string { + return fmt.Sprintf("input required: %d request(s)", len(e.InputRequests)) +} +``` + +On the server: +```go +mcp.AddTool(s, tool, func(ctx context.Context, req *mcp.CallToolRequest, in MyIn) (*mcp.CallToolResult, MyOut, error) { + if !hasConfirmation(in) { + return nil, zero, &mcp.InputRequiredError{ + InputRequests: mcp.InputRequestMap{"confirm": &mcp.ElicitParams{Message: "Sure?"}}, + RequestState: "state-token", + } + } + return &mcp.CallToolResult{ Content: []mcp.Content{&mcp.TextContent{Text: "done"}}}, myOut, nil +}) +``` + +On the client: +```go +result, err := client.CallTool(ctx, &mcp.CallToolParams{Name: "my-tool"}) +var inputReqErr *mcp.InputRequiredError +if errors.As(err, &inputReqErr) { ... } +``` + +The downsides of this approach are: +* The drift from the protocol, where MRTR is not an error flow. +* Obscure "customError -> non-error protocol type on wirte -> customError" data lifecycle. +* Things get confusing for error-processing middleware. + +--- + +### New functions + +We could introduce new functions with a different handler signature where the return type is a sealed interface. This would give us compiler-enforced correctness for values constructed by tool handlers and clients would be forced to unpack `mcp.RoundTripCallToolResult` and make a concious decision for how to handle it. +```go +type RoundTripToolHandler func(context.Context, *CallToolRequest) (RoundTripCallToolResult, error) +type RoundTripToolHandlerFor[In, Out any] func(context.Context, *CallToolRequest, In) (RoundTripCallToolResult, Out, error) + +// RoundTripCallToolResult is implemented by CallToolResult and IncompleteResult +type RoundTripCallToolResult interface { isMRTRResult() } + +type IncompleteResult struct { + ... + InputRequests InputRequestMap `json:"inputRequests,omitempty"` + RequestState string `json:"requestState,omitempty"` +} + +func (s *Server) AddRoundTripTool(t *Tool, h RoundTripToolHandler) +func AddRoundTripTool[In, Out any](s *Server, t *Tool, h RoundTripToolHandlerFor[In, Out]) +``` + +`Server.AddTool` wraps the old `ToolHandler` into a `RoundTripToolHandler` to update its function signature: +```go +mcp.AddRoundTripTool(s, tool, func(ctx context.Context, req *mcp.CallToolRequest, in MyIn) (mcp.RoundTripCallToolResult, MyOut, error) { + if needsInput(in) { + return &mcp.IncompleteResult{ + ResultType: mcp.ResultTypeInputRequired, + InputRequests: InputRequestMap{"confirm": &mcp.ElicitParams{Message: "Sure?"}}, + }, zero, nil + } + return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "done"}}}, myOut, nil +}) +``` + +The downsides of this approach are: +* SEP suggests `ResultType` will potentially be extended with new values, `RoundTrip` in new function names will not allow us to cleanly extend the sealed interface with new types. But an overly generic name for new functions will make the API use-case less clear. +* Different code +* SDK takes the same action (puts it on the wire) regardless of the returned type, it exists only for enforcing correctness of the user code. +* Exported API surface bloat: +7 exported functions. + +--- + +### Exported Middleware + +We could flip "unexported MRTR middleware with opt-out option" to "exported middleware with opt-in requirement". +```go +func AutoMRTR(opts *MultiRoundTripOptions) Middleware { ... } +type MultiRoundTripOptions struct { + MaxRetries int +} +client := mcp.NewClient(impl, nil) +client.AddSendingMiddleware(mcp.AutoMRTR(&mcp.MultiRoundTripOptions{ + MaxRetries: 5, +})) +``` +This would change semantics of `*Handler` fields - depending on the protocol version in use, an extra initialization step will be required for them to "take effect". + +--- + +### Server API protocol version bridging + +Converting `ServerSession.Elicit`/`CreateMessage`/`ListRoots` calls into MRTR wire format transparently (suspend the handler, return `input_required`, resume on retry). Rejected because of a significant implementation effort and the fact that it contradicts the design goal of MRTR where servers shouldn't hold resources between round trips, and it should be possible for a retry to arrive on any server instance in a multi-server deployment. + diff --git a/docs/CNAME b/docs/CNAME new file mode 100644 index 0000000..bd97e60 --- /dev/null +++ b/docs/CNAME @@ -0,0 +1 @@ +go.sdk.modelcontextprotocol.io \ No newline at end of file diff --git a/docs/README.md b/docs/README.md index 9f32507..d07cfa0 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,5 +1,22 @@ -# Features +# Overview + +The SDK consists of several importable packages: + +- The + [`github.com/modelcontextprotocol/go-sdk/mcp`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp) + package defines the primary APIs for constructing and using MCP clients and + servers. +- The + [`github.com/modelcontextprotocol/go-sdk/jsonrpc`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/jsonrpc) package is for users implementing + their own transports. +- The + [`github.com/modelcontextprotocol/go-sdk/auth`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth) + package provides some primitives for supporting OAuth. +- The + [`github.com/modelcontextprotocol/go-sdk/oauthex`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/oauthex) + package provides extensions to the OAuth protocol, such as ProtectedResourceMetadata. + These docs mirror the [official MCP spec](https://modelcontextprotocol.io/specification/2025-06-18). Use the index below to learn how the SDK implements a particular aspect of the @@ -38,3 +55,14 @@ protocol. # TroubleShooting See [troubleshooting.md](troubleshooting.md) for a troubleshooting guide. + +# Backwards compatibility + +See [mcpgodebug.md](mcpgodebug.md) for a list of backwards incompatible behavior changes +and description how they can be temporarily undone. + +# Rough edges + +See [rough_edges.md](rough_edges.md) for a list of rough edges or API +oversights that can't be addressed due to our compatibility promise. We'll +revisit these if/when we move to a v2 of the SDK. diff --git a/docs/client.md b/docs/client.md index 8fc7544..23d57be 100644 --- a/docs/client.md +++ b/docs/client.md @@ -2,11 +2,26 @@ # Support for MCP client features 1. [Roots](#roots) + 1. [Roots list changed](#roots-list-changed) 1. [Sampling](#sampling) 1. [Elicitation](#elicitation) + 1. [Schema defaults and enums](#schema-defaults-and-enums) + 1. [Completing a URL elicitation](#completing-a-url-elicitation) +1. [Multi Round-Trip Requests](#multi-round-trip-requests) +1. [Capabilities](#capabilities) + 1. [Capability inference](#capability-inference) + 1. [Explicit capabilities](#explicit-capabilities) + 1. [Extensions](#extensions) ## Roots +> **Note:** The roots feature is deprecated as of protocol version 2026-07-28 +> ([SEP-2577](https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging)). +> It remains fully functional during the deprecation window (at least twelve +> months). The SDK continues to support roots for compatibility. New code +> should pass paths via tool parameters, resource URIs, or configuration +> instead. + MCP allows clients to specify a set of filesystem ["roots"](https://modelcontextprotocol.io/specification/2025-06-18/client/roots). The SDK supports this as follows: @@ -24,38 +39,45 @@ client, a call to `AddRoot` or `RemoveRoots` will result in a [`ServerSession.ListRoots`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession.ListRoots) method. To receive notifications about root changes, set [`ServerOptions.RootsListChangedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.RootsListChangedHandler). +For protocol versions `2026-07-28` and later, `ListRoots` requests are +delivered via the +[Multi Round-Trip Requests](#multi-round-trip-requests) +pattern. ```go func Example_roots() { ctx := context.Background() - // Create a client with a single root. + // Create a client with two roots. c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) - c.AddRoots(&mcp.Root{URI: "file://a"}) + c.AddRoots(&mcp.Root{URI: "file://a"}, &mcp.Root{URI: "file://b"}) - // Now create a server with a handler to receive notifications about roots. - rootsChanged := make(chan struct{}) - handleRootsChanged := func(ctx context.Context, req *mcp.RootsListChangedRequest) { - rootList, err := req.Session.ListRoots(ctx, nil) - if err != nil { - log.Fatal(err) + // Create a server with a tool that requests roots via the multi round-trip + // pattern (SEP-2322): server-to-client requests are no longer sent as + // standalone JSON-RPC calls on protocol version >= 2026-07-28. + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(s, &mcp.Tool{Name: "roots"}, func(_ context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{"roots": &mcp.ListRootsParams{}}, + }, nil, nil } + rootList := req.Params.InputResponses["roots"].(*mcp.ListRootsResult) var roots []string for _, root := range rootList.Roots { roots = append(roots, root.URI) } fmt.Println(roots) - close(rootsChanged) - } - s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, &mcp.ServerOptions{ - RootsListChangedHandler: handleRootsChanged, + return &mcp.CallToolResult{}, nil, nil }) // Connect the server and client... t1, t2 := mcp.NewInMemoryTransports() - if _, err := s.Connect(ctx, t1, nil); err != nil { + serverSession, err := s.Connect(ctx, t1, nil) + if err != nil { log.Fatal(err) } + defer serverSession.Close() clientSession, err := c.Connect(ctx, t2, nil) if err != nil { @@ -63,15 +85,85 @@ func Example_roots() { } defer clientSession.Close() - // ...and add a root. The server is notified about the change. - c.AddRoots(&mcp.Root{URI: "file://b"}) - <-rootsChanged + // ...and call the tool. The client's multi round-trip driver fulfils the + // embedded roots/list request and retries the call. + if _, err := clientSession.CallTool(ctx, &mcp.CallToolParams{Name: "roots"}); err != nil { + log.Fatal(err) + } // Output: [file://a file://b] } ``` +### Roots list changed + +`Client.AddRoots` and `Client.RemoveRoots` notify every connected server that +the list changed. Servers observe this through +[`ServerOptions.RootsListChangedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.RootsListChangedHandler); +as with the server-side list-changed notifications, it reports only that +something changed, so read the list back with `ServerSession.ListRoots`. + +```go +func Example_rootsListChanged() { + ctx := context.Background() + + changed := make(chan struct{}, 2) + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, &mcp.ServerOptions{ + RootsListChangedHandler: func(context.Context, *mcp.RootsListChangedRequest) { + changed <- struct{}{} + }, + }) + + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) + c.AddRoots(&mcp.Root{URI: "file:///project"}) + + t1, t2 := mcp.NewInMemoryTransports() + ss, err := s.Connect(ctx, t1, nil) + if err != nil { + log.Fatal(err) + } + defer ss.Close() + + // ListRoots is a server-initiated request, so this session negotiates a + // protocol version that still allows one. + cs, err := c.Connect(ctx, t2, &mcp.ClientSessionOptions{ProtocolVersion: "2025-11-25"}) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + // Roots added after the client connects notify every connected server. + c.AddRoots(&mcp.Root{URI: "file:///scratch"}) + <-changed + + // The notification says only that the list changed, so read it back. + res, err := ss.ListRoots(ctx, nil) + if err != nil { + log.Fatal(err) + } + for _, root := range res.Roots { + fmt.Println(root.URI) + } + + c.RemoveRoots("file:///scratch") + <-changed + fmt.Println("roots changed again") + + // Output: + // file:///project + // file:///scratch + // roots changed again +} +``` + ## Sampling +> **Note:** The sampling feature is deprecated as of protocol version +> 2026-07-28 +> ([SEP-2577](https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging)). +> It remains fully functional during the deprecation window (at least twelve +> months). The SDK continues to support sampling for compatibility. Servers +> that need LLM completions should call LLM provider APIs directly. + [Sampling](https://modelcontextprotocol.io/specification/2025-06-18/client/sampling) is a way for servers to leverage the client's AI capabilities. It is implemented in the SDK as follows: @@ -83,6 +175,11 @@ This function is invoked whenever the server requests sampling. **Server-side**: To use sampling from the server, call [`ServerSession.CreateMessage`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession.CreateMessage). +For protocol versions `2026-07-28` and later, sampling requests are +delivered via the +[Multi Round-Trip Requests](#multi-round-trip-requests) +pattern. + ```go func Example_sampling() { ctx := context.Background() @@ -100,22 +197,35 @@ func Example_sampling() { // Connect the server and client... ct, st := mcp.NewInMemoryTransports() + // Create a server with a tool that requests sampling via the multi + // round-trip pattern (SEP-2322): server-to-client requests are no longer + // sent as standalone JSON-RPC calls on protocol version >= 2026-07-28. s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(s, &mcp.Tool{Name: "sample"}, func(_ context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{"msg": &mcp.CreateMessageParams{}}, + }, nil, nil + } + msg := req.Params.InputResponses["msg"].(*mcp.CreateMessageWithToolsResult) + return &mcp.CallToolResult{Content: msg.Content}, nil, nil + }) session, err := s.Connect(ctx, st, nil) if err != nil { log.Fatal(err) } defer session.Close() - if _, err := c.Connect(ctx, ct, nil); err != nil { - log.Fatal(err) - } - - msg, err := session.CreateMessage(ctx, &mcp.CreateMessageParams{}) + clientSession, err := c.Connect(ctx, ct, nil) if err != nil { log.Fatal(err) } - fmt.Println(msg.Content.(*mcp.TextContent).Text) + + res, err := clientSession.CallTool(ctx, &mcp.CallToolParams{Name: "sample"}) + if err != nil { + log.Fatal(err) + } + fmt.Println(res.Content[0].(*mcp.TextContent).Text) // Output: would have created a message } ``` @@ -128,17 +238,45 @@ allows servers to request user inputs. It is implemented in the SDK as follows: **Client-side**: To add the `elicitation` capability to a client, set [`ClientOptions.ElicitationHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.ElicitationHandler). The elicitation handler must return a result that matches the requested schema; -otherwise, elicitation returns an error. +otherwise, elicitation returns an error. If your handler supports [URL mode +elicitation](https://modelcontextprotocol.io/specification/2025-11-25/client/elicitation#url-mode-elicitation-requests), +you must declare that capability explicitly (see [Capabilities](#capabilities)) **Server-side**: To use elicitation from the server, call [`ServerSession.Elicit`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession.Elicit). +For protocol versions `2026-07-28` and later, elicitation requests are +delivered via the +[Multi Round-Trip Requests](#multi-round-trip-requests) +pattern. + ```go func Example_elicitation() { ctx := context.Background() ct, st := mcp.NewInMemoryTransports() + // Create a server with a tool that requests elicitation via the multi + // round-trip pattern (SEP-2322): server-to-client requests are no longer + // sent as standalone JSON-RPC calls on protocol version >= 2026-07-28. s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(s, &mcp.Tool{Name: "ask"}, func(_ context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{"input": &mcp.ElicitParams{ + Message: "This should fail", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "test": {Type: "string"}, + }, + }, + }}, + }, nil, nil + } + res := req.Params.InputResponses["input"].(*mcp.ElicitResult) + fmt.Println(res.Content["test"]) + return &mcp.CallToolResult{}, nil, nil + }) ss, err := s.Connect(ctx, st, nil) if err != nil { log.Fatal(err) @@ -150,22 +288,261 @@ func Example_elicitation() { return &mcp.ElicitResult{Action: "accept", Content: map[string]any{"test": "value"}}, nil }, }) - if _, err := c.Connect(ctx, ct, nil); err != nil { - log.Fatal(err) - } - res, err := ss.Elicit(ctx, &mcp.ElicitParams{ - Message: "This should fail", - RequestedSchema: &jsonschema.Schema{ - Type: "object", - Properties: map[string]*jsonschema.Schema{ - "test": {Type: "string"}, - }, - }, - }) + clientSession, err := c.Connect(ctx, ct, nil) if err != nil { log.Fatal(err) } - fmt.Println(res.Content["test"]) + if _, err := clientSession.CallTool(ctx, &mcp.CallToolParams{Name: "ask"}); err != nil { + log.Fatal(err) + } // Output: value } ``` + +### Schema defaults and enums + +`ElicitParams.RequestedSchema` is a flat schema of primitive fields, which the +client renders as a form. Two field keywords shape that form. + +A `Default` ([SEP-1034](https://modelcontextprotocol.io/seps/1034--support-default-values-for-all-primitive-types-in)) prefills a +field. When the user accepts without supplying it, the SDK fills the field in +from the schema before the result reaches either side's caller — the client +does so after its elicitation handler returns, and `ServerSession.Elicit` does +so again on receipt. This is unconditional; there is no opt-in flag. Marking a +defaulted field `Required` defeats it: accepted content is validated against +the schema before defaults are applied, so an answer that omits the field is +rejected rather than defaulted. + +An `Enum` ([SEP-1330](https://modelcontextprotocol.io/seps/1330-elicitation-enum-schema-improvements-and-standards)) restricts a +field to a fixed set of values, which the client renders as a choice. Enums +are supported only on `"string"` fields; declaring one on another type is +rejected. To label the choices, set the legacy `enumNames` keyword through +`Schema.Extra`, with exactly one name per enum value — a mismatched length is +rejected. + +```go +func Example_elicitationSchema() { + ctx := context.Background() + ct, st := mcp.NewInMemoryTransports() + + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(s, &mcp.Tool{Name: "export_report"}, func(_ context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{"format": &mcp.ElicitParams{ + Message: "Export quarterly-sales as which format?", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "format": { + Type: "string", + Title: "Format", + Enum: []any{"pdf", "csv"}, + Default: json.RawMessage(`"pdf"`), + Extra: map[string]any{"enumNames": []any{"PDF document", "CSV spreadsheet"}}, + }, + }, + }, + }}, + }, nil, nil + } + res := req.Params.InputResponses["format"].(*mcp.ElicitResult) + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "Exported as " + res.Content["format"].(string)}}, + }, nil, nil + }) + if _, err := s.Connect(ctx, st, nil); err != nil { + log.Fatal(err) + } + + // The user accepts without filling anything in. + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, &mcp.ClientOptions{ + ElicitationHandler: func(context.Context, *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + return &mcp.ElicitResult{Action: "accept", Content: map[string]any{}}, nil + }, + }) + cs, err := c.Connect(ctx, ct, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + res, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: "export_report"}) + if err != nil { + log.Fatal(err) + } + fmt.Println(res.Content[0].(*mcp.TextContent).Text) + // Output: Exported as pdf +} +``` + +### Completing a URL elicitation + +In URL mode the user finishes out of band, in a browser, so nothing in the +elicitation result tells the client when they are done. The server signals that +with +[`ServerSession.NotifyElicitationComplete`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession.NotifyElicitationComplete), +passing the same `ElicitationID` the request carried; the client observes it +through +[`ClientOptions.ElicitationCompleteHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.ElicitationCompleteHandler). +Send it from whatever endpoint the hosted flow redirects back to. + +The notification matters most when a handler rejects a request with +[`URLElicitationRequiredError`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#URLElicitationRequiredError): +the client parks the original request until a notification names that +`ElicitationID`, and then retries it automatically. Until one arrives, the +client waits. + +```go +func Example_elicitationComplete() { + ctx := context.Background() + ct, st := mcp.NewInMemoryTransports() + + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + log.Fatal(err) + } + defer ss.Close() + + done := make(chan struct{}) + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, &mcp.ClientOptions{ + Capabilities: &mcp.ClientCapabilities{ + Elicitation: &mcp.ElicitationCapabilities{URL: &mcp.URLElicitationCapabilities{}}, + }, + ElicitationHandler: func(_ context.Context, req *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + fmt.Println("opening", req.Params.URL) + return &mcp.ElicitResult{Action: "accept"}, nil + }, + ElicitationCompleteHandler: func(_ context.Context, req *mcp.ElicitationCompleteNotificationRequest) { + fmt.Println("flow finished:", req.Params.ElicitationID) + close(done) + }, + }) + cs, err := c.Connect(ctx, ct, &mcp.ClientSessionOptions{ProtocolVersion: "2025-11-25"}) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + const elicitationID = "connect-calendar-1" + if _, err := ss.Elicit(ctx, &mcp.ElicitParams{ + Message: "Grant calendar access", + URL: "https://calendar.example.com/consent?state=" + elicitationID, + ElicitationID: elicitationID, + }); err != nil { + log.Fatal(err) + } + + // The hosted page redirects back to the server, whose callback endpoint + // signals that the user is done. + if err := ss.NotifyElicitationComplete(ctx, &mcp.ElicitationCompleteParams{ElicitationID: elicitationID}); err != nil { + log.Fatal(err) + } + <-done + + // Output: + // opening https://calendar.example.com/consent?state=connect-calendar-1 + // flow finished: connect-calendar-1 +} +``` + +## Multi Round-Trip Requests + +[SEP-2322](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2322) +introduces the MRTR pattern: server-to-client requests for sampling, +elicitation, and roots are no longer issued as fresh JSON-RPC requests but +are carried inside the in-flight reply of a `tools/call`, `prompts/get`, or +`resources/read`. The client must respond by retrying the original request +with the produced responses. + +The SDK installs `clientMultiRoundTripMiddleware` for every client by +default. The middleware: + +1. Inspects each `tools/call`/`prompts/get`/`resources/read` reply. +2. If the result's `NeedsInput()` is true, fans out the `InputRequests` map + concurrently, calling the configured handler for each (`elicit`, + `createMessage`/`createMessageWithTools`, or `listRoots`). +3. Threads the server-supplied opaque `RequestState` back unchanged. +4. Retries the original request with the responses set, repeating until the + result no longer needs input. + +The middleware is enabled by default. To opt out, set +[`ClientOptions.MultiRoundTrip.Disabled = true`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#MultiRoundTripOptions); +the client will then surface input-required results directly to the caller +(the returned `CallToolResult`, `GetPromptResult`, or `ReadResourceResult` +will report `NeedsInput() == true` and expose the server's `InputRequests` +and opaque `RequestState`). Your code must fulfil each request and re-issue +the original call with `InputResponses` set and `RequestState` echoed +back. + +For legacy (`<= 2025-11-25`) servers, the SDK transparently sends server +requests on the legacy server-initiated channel; the MRTR machinery is a +no-op in that direction. For legacy clients talking to MRTR-style servers, +the server SDK applies the inverse compatibility shim — see the +[server-side documentation](server.md#multi-round-trip-requests). + +## Capabilities + +Client capabilities are advertised to servers during the initialization +handshake. [By default](rough_edges.md), the SDK advertises the `logging` +capability. Additional capabilities are automatically added when server +features are added (e.g. via `AddTool`), or when handlers are set in the +`ServerOptions` struct (e.g., setting `CompletionHandler` adds the +`completions` capability), or may be configured explicitly. + +### Capability inference + +When handlers are set on `ClientOptions` (e.g., `CreateMessageHandler` or +`ElicitationHandler`), the corresponding capability is automatically added if +not already present, with a default configuration. + +For elicitation, if the handler is set but no `Capabilities.Elicitation` is +specified, the client defaults to form elicitation. To enable URL elicitation +or both modes, [configure `Capabilities.Elicitation` +explicitly](#explicit-capabilities). + +See the [`ClientCapabilities` +documentation](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientCapabilities) +for further details on inference. + +### Explicit capabilities + +To explicitly declare capabilities, or to override the [default inferred +capability](#capability-inference), set +[`ClientOptions.Capabilities`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.Capabilities). +This sets the initial client capabilities, before any capabilities are added +based on configured handlers. If a capability is already present in +`Capabilities`, adding a handler will not change its configuration. + +This allows you to: + +- **Disable default capabilities**: Pass an empty `&ClientCapabilities{}` to + disable all defaults, including roots. +- **Disable listChanged notifications**: Set `ListChanged: false` on a + capability to prevent the client from sending list-changed notifications + when roots are added or removed. +- **Configure elicitation modes**: Specify which elicitation modes (form, URL) + the client supports. + +```go +// Configure elicitation modes and disable roots. +client := mcp.NewClient(impl, &mcp.ClientOptions{ + Capabilities: &mcp.ClientCapabilities{ + Elicitation: &mcp.ElicitationCapabilities{ + Form: &mcp.FormElicitationCapabilities{}, + URL: &mcp.URLElicitationCapabilities{}, + }, + }, + ElicitationHandler: handler, +}) +``` + +### Extensions + +[SEP-2133](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2133) +adds an `extensions` map to `ClientCapabilities` and `ServerCapabilities` so +that optional capabilities outside the core protocol can be declared on the +wire. Keys are namespaced as `"{vendor-prefix}/{extension-name}"`; values +are per-extension settings objects. + diff --git a/docs/contributing.md b/docs/contributing.md new file mode 100644 index 0000000..ea38c9b --- /dev/null +++ b/docs/contributing.md @@ -0,0 +1 @@ +--8<-- "CONTRIBUTING.md" diff --git a/docs/images/logo-dark.svg b/docs/images/logo-dark.svg new file mode 100644 index 0000000..03d9f85 --- /dev/null +++ b/docs/images/logo-dark.svg @@ -0,0 +1,12 @@ + + + + + + + + + + + + diff --git a/docs/images/logo-light.svg b/docs/images/logo-light.svg new file mode 100644 index 0000000..fe5edb7 --- /dev/null +++ b/docs/images/logo-light.svg @@ -0,0 +1,69 @@ + + + + + + + + + + + + + + diff --git a/docs/mcpgodebug.md b/docs/mcpgodebug.md new file mode 100644 index 0000000..43a62ce --- /dev/null +++ b/docs/mcpgodebug.md @@ -0,0 +1,178 @@ + +# Backwards compatibility and MCPGODEBUG + + According to our compatibility promise, we can't break backward compatibility + of the SDK API. However, sometimes we need to change the behavior of the SDK + in a backward-incompatible way in order to fix bugs or security issues. + In those cases we introduce temporary compatibility parameters, that can be + used to opt-out of the new behavior. They are usually maintained for two + minor release cycles and then removed. + + The compatibility parameters are provided via the `MCPGODEBUG` environment + variable. The value of the variable is a comma-separated list of parameter + value assignments, e.g.: + + ``` + MCPGODEBUG=parameter1=value1,parameter2=value2 + ``` + +## `MCPGODEBUG` history + +### 1.9.0 + +Options listed below were added and will be removed in the 1.11.0 version of the SDK. + +- `structuredcontentfloat64` added. If set to `1`, numbers in + `CallToolResult.StructuredContent` are decoded as `float64`, restoring the + previous behavior. The default behavior decodes numbers as `json.Number` so + that large integers and other exact JSON values do not silently lose + precision. + +### 1.8.0 + +Options listed below were added and will be removed in the 1.9.0 version of the SDK. + +- `plaintextstatefulrejection` added. If set to `1`, a stateful + `StreamableHTTPHandler` will respond with a plain-text `http.Error` 400 body + when it receives a request carrying per-request metadata (i.e. an + `io.modelcontextprotocol/protocolVersion` `_meta` field, or an + `MCP-Protocol-Version` header >= `2026-07-28`), restoring the previous + behavior. The default behavior was changed so that the server responds with a + JSON-RPC error of code `CodeUnsupportedProtocolVersion` (`-32022`) carrying + an `UnsupportedProtocolVersionData` payload that advertises the legacy + versions the server supports. This lets the client's + existing renegotiation logic recover and prevents the failure from tearing + down the underlying connection. + +- `blockingcancelnotify` added. If set to `1`, a cancelled call waits + synchronously for the best-effort `notifications/cancelled` message to be + delivered (up to `notifyCancellationTimeout`, currently 5 s) before + returning to the caller, restoring the previous behavior. The delivery + error is joined into the caller's returned error. The default behavior was + changed so that the call is retired immediately and the notification is + sent asynchronously off the caller's return path: the caller returns as + soon as its context is cancelled and cannot be delayed by a slow or + unresponsive peer. See issue #1150. + +Options below were removed, according to plan: + +- `seterroroverwrite`. `SetError` now always preserves existing `Content`. + +- `enableoriginverification`. A nil + `StreamableHTTPOptions.CrossOriginProtection` now always means no + cross-origin protection. Wrap the handler with cross-origin protection + middleware instead. + +- `disablecontenttypecheck`. Content-Type validation on HTTP POST requests can + no longer be disabled. + +- `disablelocalhostprotection`, whose removal had been postponed to this + release. Set the `DisableLocalhostProtection` field in + `StreamableHTTPOptions` or `SSEOptions` to opt out of DNS rebinding + protection. + +### 1.7.0 + +Options listed below were added and will be removed in the 1.9.0 version of the SDK. + +- `customresnotfounderrcode` added. If set to `1`, `ResourceNotFoundError` will + use the custom error code `-32002` instead of the standard `-32602` (Invalid + Params), restoring the previous behavior. The default behavior was changed to + align with SEP-2164 and the JSON-RPC specification. + +- `hintomitempty` added. If set to `1`, `ToolAnnotations` JSON marshaling + will omit `ReadOnlyHint` and `IdempotentHint` when their value is `false`, + restoring the previous behavior. The default behavior was changed to always + serialize these fields, since their Go types are bare `bool` (not `*bool`) + and omitting `false` made it indistinguishable from unset. + +- `allowsessionsinstateless` added. If set to `1`, stateless streamable HTTP + servers will read the `Mcp-Session-Id` request header (or generate one via + `GetSessionID`), set it on response headers, and accept `DELETE` requests, + restoring the previous behavior. The default behavior was changed so that + stateless servers ignore session IDs entirely and reject `DELETE` with 405. + +- `nomethodnotfoundcodeinerror` added. If set to `1`, the jsonrpc2 layer will not + include the MethodNotFound Error (`-32601`) in the error response when the + requested method in STDIO transport is not found. The default behavior was + changed to include the MethodNotFound Error in the error response when the + requested method in STDIO transport is not found. + +- `noprotocolerrorbody` added. If set to `1`, the streamable HTTP client will + not attempt to decode the JSON-RPC error body of a non-2xx HTTP response, + and any non-transient error will permanently fail the connection, restoring + the previous behavior. The default behavior was changed so that the client + attempts to decode the JSON-RPC error body of a non-2xx response, surfaces + the underlying JSON-RPC error, and wraps it with `jsonrpc2.ErrRejected` so + that per-call rejections do not tear down the session. + +- `nowrapinvalidparams` added. If set to `1`, the server will not wrap + params-decoding failures with `jsonrpc2.ErrInvalidParams`, so wire responses + carry the zero-value error code `0` instead of `-32602` ("invalid params"), + restoring the previous behavior. The default behavior was changed so that + the server always reports params-decoding failures with the standard + `-32602` code. + +- `disablecompleteparamsvalidation` added. If set to `1`, `Server.complete` + will not validate that the required `ref` and `argument.name` fields on + `CompleteParams` are present, restoring the previous behavior of dispatching + the request to the completion handler unconditionally. The default behavior + was changed to reject malformed requests with `-32602` (Invalid Params). + +### 1.6.1 + +Options listed below were added and will be removed in the 1.8.0 version of the SDK. + +- `disablecontenttypecheck` added. If set to `1`, Content-Type validation on + HTTP POST requests will be disabled, allowing requests with non-JSON or missing + Content-Type headers. The default behavior is to validate that HTTP POST + requests have Content-Type: application/json. + +### 1.6.0 + +Options listed below were added and will be removed in the 1.8.0 version of the SDK. + +- `seterroroverwrite` added. If set to `1`, `SetError` will always overwrite + `Content` with the error text, restoring the previous behavior. The default + behavior was changed to preserve existing `Content` if it has already been + populated. + +- `enableoriginverification` added. If set to `1`, default (zero-value) + cross-origin protection will be applied when + `StreamableHTTPOptions.CrossOriginProtection` is nil, restoring the + behavior from v1.4.1-v1.5.0. The default behavior was changed to not + enable cross-origin protection. + +- `disablelocalhostprotection` removal was postponed until 1.8.0, as it is now + also used in the SSE transport. + +Options below were removed: + +- `jsonescaping`, according to plan, + +- `disablecrossoriginprotection`, it was replaced by + `enableoriginverification` after the default was changed to not enable + cross-origin protection. + +### 1.4.1 + +Options listed below will be removed in the 1.6.0 version of the SDK. + +- `disablecrossoriginprotection` added. If set to `1`, newly added cross-origin + protection will be disabled. The default behavior was changed to enable + cross-origin protection. + +### 1.4.0 + +Options listed below will be removed in the 1.6.0 version of the SDK. + +- `jsonescaping` added. If set to `1`, JSON marshaling will preserve the previous + behavior of escaping HTML characters in JSON strings. The default behavior + was changed to not escape HTML characters, to be consistent with other SDKs. + +- `disablelocalhostprotection` added. If set to `1`, newly added DNS rebinding + protection will be disabled. The default behavior was changed to enable DNS rebinding + protection. The protection can also be disabled by setting the + `DisableLocalhostProtection` field in the `StreamableHTTPOptions` or + `SSEOptions` struct to `true`, which is the recommended way to disable + the protection long term. **Removal of this option was postponed until 1.8.0.** diff --git a/docs/protocol.md b/docs/protocol.md index 3a031af..06c1b77 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -2,9 +2,14 @@ # Support for the MCP base protocol 1. [Lifecycle](#lifecycle) + 1. [Discovery](#discovery) + 1. [Per-request metadata keys](#per-request-metadata-keys) + 1. [Per-response metadata keys](#per-response-metadata-keys) + 1. [Subscriptions](#subscriptions) 1. [Transports](#transports) 1. [Stdio Transport](#stdio-transport) 1. [Streamable Transport](#streamable-transport) + 1. [Legacy SSE Transport](#legacy-sse-transport) 1. [Custom transports](#custom-transports) 1. [Concurrency](#concurrency) 1. [Authorization](#authorization) @@ -13,25 +18,36 @@ 1. [Security](#security) 1. [Confused Deputy](#confused-deputy) 1. [Token Passthrough](#token-passthrough) + 1. [Server-Side Request Forgery](#server-side-request-forgery) 1. [Session Hijacking](#session-hijacking) + 1. [Issuer Mix-Up](#issuer-mix-up) 1. [Utilities](#utilities) 1. [Cancellation](#cancellation) 1. [Ping](#ping) 1. [Progress](#progress) + 1. [Error codes](#error-codes) ## Lifecycle The SDK provides an API for defining both MCP clients and servers, and connecting them over various transports. When a client and server are -connected, it creates a logical session, which follows the MCP spec's -[lifecycle](https://modelcontextprotocol.io/specification/2025-06-18/basic/lifecycle). +connected, it creates a logical session. -In this SDK, both a -[`Client`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Client) -and -[`Server`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server) -can handle multiple peers. Every time a new peer is connected, it creates a new -session. +The MCP specification defines two lifecycle models, and the SDK supports both +transparently based on the negotiated protocol version: + +- The **legacy `initialize` handshake** + ([lifecycle](https://modelcontextprotocol.io/specification/2025-06-18/basic/lifecycle)) + used by protocol versions through `2025-11-25`. +- A **stateless** model introduced in `2026-07-28` by + [SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575), + in which there is no `initialize`/`notifications/initialized` handshake, and + each request carries its protocol version and client capabilities in + `_meta`. Clients SHOULD also include their identity (`clientInfo`) on every + request, and servers SHOULD include their identity (`serverInfo`) on every + response. + +In both models, the SDK exposes the same API: - A `Client` is a logical MCP client, configured with various [`ClientOptions`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions). @@ -47,11 +63,16 @@ session. [`Server.Connect`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server.Connect), it creates a [`ServerSession`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession). - This session is not initialized until the client sends the - `notifications/initialized` message. Use `ServerOptions.InitializedHandler` - to listen for this event, or just use the session through various feature - handlers (such as a `ToolHandler`). Requests to the server are rejected until - the client has initialized the session. + +In the legacy model, the server session is not considered initialized until +the client sends the `notifications/initialized` message. Use +`ServerOptions.InitializedHandler` to listen for this event, or just use the +session through various feature handlers (such as a `ToolHandler`). Requests +to the server are rejected until the client has initialized the session. + +In the stateless model (`2026-07-28`+), there is no handshake. The server +processes the first request the moment it arrives, validates the per-request +`_meta` fields. Both `ClientSession` and `ServerSession` have a `Close` method to terminate the session, and a `Wait` method to await session termination by the peer. Typically, @@ -61,19 +82,13 @@ it is the client's responsibility to end the session. func Example_lifecycle() { ctx := context.Background() - // Create a client and server. - // Wait for the client to initialize the session. + // Create a client and server. Under protocol version 2026-07-28, there + // is no dedicated initialize/initialized handshake: the session is + // implicitly live the moment the client issues its first request. client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) - server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, &mcp.ServerOptions{ - InitializedHandler: func(context.Context, *mcp.InitializedRequest) { - fmt.Println("initialized!") - }, - }) + server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) // Connect the server and client using in-memory transports. - // - // Connect the server first so that it's ready to receive initialization - // messages from the client. t1, t2 := mcp.NewInMemoryTransports() serverSession, err := server.Connect(ctx, t1, nil) if err != nil { @@ -92,10 +107,86 @@ func Example_lifecycle() { if err := serverSession.Wait(); err != nil { log.Fatal(err) } - // Output: initialized! + // Output: } ``` +### Discovery + +Introduced in `2026-07-28` by +[SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575), +the `server/discover` RPC lets a client discover a server's supported +protocol versions, capabilities, and identity before issuing any other +request. Servers implementing `2026-07-28` MUST implement it. + +- **Server**: `Server.Connect` registers the + `server/discover` handler automatically; the response is computed from the + server's static configuration (capabilities, server info, instructions) and + the transport-filtered list of supported protocol versions. +- **Client**: `Client.Connect` calls `server/discover` first and + uses the result to negotiate a mutually supported version. If discovery + fails or the server does not support the latest version, the client falls back to the + legacy `initialize` handshake. + +A server advertises and negotiates every protocol version the SDK supports; +set [`ServerOptions.SupportedProtocolVersions`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions) +to narrow that set, for example to stop serving a revision your deployment has +retired. + +```go +server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v1.0.0"}, &mcp.ServerOptions{ + SupportedProtocolVersions: []string{"2026-07-28", "2025-11-25"}, +}) +``` + +The set can only narrow support, never widen it; the field's documentation +covers how each protocol era answers a request at an excluded version. + +### Per-request metadata keys + +When the negotiated protocol version is `2026-07-28` or later, every request +carries these keys inside its `_meta` map (constants live in +[`mcp/protocol.go`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#pkg-constants)): + +| Constant | Wire key | Type | Required | +|---|---|---|---| +| `MetaKeyProtocolVersion` | `io.modelcontextprotocol/protocolVersion` | `string` | Yes | +| `MetaKeyClientCapabilities` | `io.modelcontextprotocol/clientCapabilities` | `*ClientCapabilities` | Yes | +| `MetaKeyClientInfo` | `io.modelcontextprotocol/clientInfo` | `*Implementation` | No | +| `MetaKeyLogLevel` | `io.modelcontextprotocol/logLevel` | `LoggingLevel` (deprecated by SEP-2577) | No | + +The client populates the required keys automatically on every outgoing +request, and populates `clientInfo` when configured with an `*Implementation` +(the default). Server-side handlers can read them through +`ServerRequest[P].ProtocolVersion()`, `ServerRequest[P].ClientInfo()`, and +`ServerRequest[P].ClientCapabilities()`. + +### Per-response metadata keys + +Under the same protocol version, servers SHOULD identify themselves on every +response. The SDK populates this +key automatically on every outgoing response: + +| Constant | Wire key | Type | Required | +|---|---|---|---| +| `MetaKeyServerInfo` | `io.modelcontextprotocol/serverInfo` | `*Implementation` | No | + +### Subscriptions + +Introduced in `2026-07-28` by +[SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575), +`subscriptions/listen` replaces the legacy `resources/subscribe` RPC and +the GET-based SSE endpoint with a single long-lived request that +multiplexes every kind of server-to-client change notification a client +opts in to. On the wire the client sends one `subscriptions/listen` +request whose `notifications` field enumerates what it wants +(`toolsListChanged`, `promptsListChanged`, `resourcesListChanged`, and/or +a list of resource URIs in `resourceSubscriptions`); the server replies +first with a `notifications/subscriptions/acknowledged` notification +reporting the honored subset, then streams every change notification on +the same request, and finally closes with a `SubscriptionsListenResult` +when it tears the subscription down. + ## Transports A @@ -148,6 +239,7 @@ To create a streamable MCP server, you create a `StreamableHTTPHandler` and pass it an `mcp.Server`: ```go +// TODO: Until we have a way to clean up abandoned sessions, this test will leak goroutines (see #499) func ExampleStreamableHTTPHandler() { // Create a new streamable handler, using the same MCP server for every request. // @@ -161,10 +253,10 @@ func ExampleStreamableHTTPHandler() { defer httpServer.Close() // The SDK is currently permissive of some missing keys in "params". - resp := mustPostMessage(`{"jsonrpc": "2.0", "id": 1, "method":"initialize", "params": {}}`, httpServer.URL) + resp := mustPostMessage(`{"jsonrpc": "2.0", "id": 1, "method":"initialize", "params": {"protocolVersion":"2025-11-25"}}`, httpServer.URL) fmt.Println(resp) // Output: - // {"jsonrpc":"2.0","id":1,"result":{"capabilities":{"logging":{}},"protocolVersion":"2025-06-18","serverInfo":{"name":"server","version":"v0.1.0"}}} + // {"jsonrpc":"2.0","id":1,"result":{"capabilities":{"logging":{}},"protocolVersion":"2025-11-25","serverInfo":{"name":"server","version":"v0.1.0"}}} } ``` @@ -179,12 +271,66 @@ connect to the server: transport := &mcp.StreamableClientTransport{ Endpoint: "http://localhost:8080/mcp", } -client, err := mcp.Connect(ctx, transport, &mcp.ClientOptions{...}) +client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v1.0.0"}, nil) +session, err := client.Connect(ctx, transport, nil) ``` The `StreamableClientTransport` handles the HTTP requests and communicates with the server using the streamable transport protocol. +#### HTTP Headers + +[SEP-2243](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2243) +standardises a small set of MCP-specific HTTP headers, sent and validated by +the SDK only for sessions negotiated on protocol version `2026-07-28` or +later. + +| Header | Direction | Purpose | +|---|---|---| +| `Mcp-Protocol-Version` | request | Mirrors `_meta.io.modelcontextprotocol/protocolVersion` from the body | +| `Mcp-Session-Id` | request/response | Logical session identifier (removed in stateless mode) | +| `Mcp-Method` | request | Mirrors the JSON-RPC `method` field; mismatch ⇒ `-32020` | +| `Mcp-Name` | request | Mirrors the request's principal name (`tools/call.params.name`, `prompts/get.params.name`, `resources/read.params.uri`); mismatch ⇒ `-32020` | +| `Mcp-Param-{Header}` | request | Per-tool parameter passthrough; see below | + +The header-name set is exhaustive: for protocol versions earlier than +`2026-07-28`, only `Mcp-Protocol-Version` and `Mcp-Session-Id` are recognised. + +**Body↔header mirroring.** When `Mcp-Method` or `Mcp-Name` is present but +does not match the JSON-RPC body, the server returns +[`CodeHeaderMismatch`](#error-codes) (`-32020`). The same applies to a +mismatch between `Mcp-Protocol-Version` and the body's +`_meta.io.modelcontextprotocol/protocolVersion`. + +**`Mcp-Param-*` passthrough.** A tool may annotate properties of its +`InputSchema` with `x-mcp-header: ""`. When a client invokes that +tool over HTTP, the SDK serializes the matching argument(s) as +`Mcp-Param-` request headers (using `=?base64?...?=` encoding for +non-ASCII values). On the server, the headers are validated against the body +and either accepted or rejected with `-32020`. Properties so annotated must +be typed `string`, `integer`, or `boolean`; tools that fail validation at +registration time are silently dropped from `tools/list`. + +#### Resumability and Redelivery + +By default, the streamable server does not support [resumability or +redelivery](https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#resumability-and-redelivery) +of messages, because doing so requires either a persistent storage solution or +unbounded memory usage (see also +[#580](https://github.com/modelcontextprotocol/go-sdk/issues/580)). + +To enable resumability, set `StreamableHTTPOptions.EventStore` to a non-nil +value. The SDK provides a `MemoryEventStore` for testing or simple use cases; +for production use it is generally advisable to use a more sophisticated +implementation. + +> **Note**: [SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575) +> removes SSE stream resumability (`Last-Event-ID`, SSE event IDs) for protocol +> version `2026-07-28`. The SDK preserves the `EventStore` and `Last-Event-ID` +> code paths for backward compatibility with `2025-11-25` and earlier; on +> `2026-07-28` sessions, a broken response stream loses the in-flight request +> and the client must re-issue it as a new request with a new ID. + #### Stateless Mode The streamable server supports a _stateless mode_ by setting @@ -203,10 +349,71 @@ to see the logical session > modelcontextprotocol/modelcontextprotocol#1372, or > modelcontextprotocol/modelcontextprotocol#1442 for potential refinements. -_See [examples/server/distributed](../examples/server/distributed/main.go) for -an example using statless mode to implement a server distributed across +> **Required for `2026-07-28`**: the streamable HTTP transport accepts +> requests at protocol version `2026-07-28` **only** when `Stateless = true`. +> Requests at that version against a non-stateless handler are rejected. + +_See [examples/server/distributed](https://github.com/modelcontextprotocol/go-sdk/blob/main/examples/server/distributed/main.go) for +an example using stateless mode to implement a server distributed across multiple processes._ +### Legacy SSE Transport + +Before the streamable transport, the +[2024-11-05](https://modelcontextprotocol.io/specification/2024-11-05/basic/transports) +spec defined an HTTP transport built from two endpoints: a hanging GET that +streams server-to-client messages as server-sent events, and a per-session +endpoint the client POSTs its messages to. The SDK still implements it, for +talking to peers that predate the streamable transport. New deployments should +use the [streamable transport](#streamable-transport) instead. + +It is implemented across three types, mirroring the streamable ones: + +- [`SSEHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#SSEHandler): + an `http.Handler` that creates a session per incoming GET and routes POSTs to + the right session. Construct it with + [`NewSSEHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#NewSSEHandler), + which takes a function returning the `Server` to serve a given request. +- [`SSEServerTransport`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#SSEServerTransport): + one such session. Reach for it directly only when serving the two endpoints + yourself; `SSEHandler` creates these for you. +- [`SSEClientTransport`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#SSEClientTransport): + the client side, which needs only the `Endpoint` of the GET. + +```go +func ExampleSSEHandler() { + server := mcp.NewServer(&mcp.Implementation{Name: "adder", Version: "v0.0.1"}, nil) + mcp.AddTool(server, &mcp.Tool{Name: "add", Description: "add two numbers"}, Add) + + handler := mcp.NewSSEHandler(func(*http.Request) *mcp.Server { return server }, nil) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + ctx := context.Background() + transport := &mcp.SSEClientTransport{Endpoint: httpServer.URL} + client := mcp.NewClient(&mcp.Implementation{Name: "test", Version: "v1.0.0"}, nil) + cs, err := client.Connect(ctx, transport, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + res, err := cs.CallTool(ctx, &mcp.CallToolParams{ + Name: "add", + Arguments: map[string]any{"x": 1, "y": 2}, + }) + if err != nil { + log.Fatal(err) + } + fmt.Println(res.Content[0].(*mcp.TextContent).Text) + + // Output: 3 +} +``` + +_See [examples/server/sse](https://github.com/modelcontextprotocol/go-sdk/blob/main/examples/server/sse/main.go) +for a standalone server._ + ### Custom transports The SDK supports [custom @@ -215,7 +422,7 @@ by implementing the [`Transport`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Transport) interface: a logical bidirectional stream of JSON-RPC messages. -_Full example: [examples/server/custom-transport](../examples/server/custom-transport/main.go)._ +_Full example: [examples/server/custom-transport](https://github.com/modelcontextprotocol/go-sdk/blob/main//examples/server/custom-transport/main.go)._ ### Concurrency @@ -249,7 +456,7 @@ and invokes the passed to `RequireBearerToken` to parse the token and perform validation. The middleware function checks expiration and scopes (if they are provided in [`RequireBearerTokenOptions.Scopes`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#RequireBearerTokenOptions.Scopes)), so the -`TokenVerifer` doesn't have to. +`TokenVerifier` doesn't have to. If [`RequireBearerTokenOptions.ResourceMetadataURL`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#RequireBearerTokenOptions.ResourceMetadataURL) is set and verification fails, the middleware function sets the WWW-Authenticate header as required by the [Protected Resource Metadata spec](https://datatracker.ietf.org/doc/html/rfc9728). @@ -259,38 +466,201 @@ from `req.Extra.TokenInfo`, where `req` is the handler's request. (For example, [`CallToolRequest`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#CallToolRequest).) HTTP handlers wrapped by the `RequireBearerToken` middleware can obtain the `TokenInfo` from the context with [`auth.TokenInfoFromContext`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#TokenInfoFromContext). - + +#### OAuth Protected Resource Metadata + +Servers implementing OAuth 2.0 authorization should expose a protected resource metadata endpoint +as specified in [RFC 9728](https://datatracker.ietf.org/doc/html/rfc9728). This endpoint provides +clients with information about the resource server's OAuth configuration, including which +authorization servers can be used and what scopes are supported. + +The SDK provides [`ProtectedResourceMetadataHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#ProtectedResourceMetadataHandler) +to serve this metadata. The handler automatically sets CORS headers (`Access-Control-Allow-Origin: *`) +to support cross-origin client discovery, as the metadata contains only public configuration information. + +Example usage: + +```go +metadata := &oauthex.ProtectedResourceMetadata{ + Resource: "https://example.com/mcp", + AuthorizationServers: []string{ + "https://auth.example.com/.well-known/openid-configuration", + }, + ScopesSupported: []string{"read", "write"}, +} +http.Handle("/.well-known/oauth-protected-resource", + auth.ProtectedResourceMetadataHandler(metadata)) +``` + +For more sophisticated CORS policies, wrap the handler with a CORS middleware like +[github.com/rs/cors](https://github.com/rs/cors) or [github.com/jub0bs/cors](https://github.com/jub0bs/cors). The [_auth middleware example_](https://github.com/modelcontextprotocol/go-sdk/tree/main/examples/server/auth-middleware) shows how to implement authorization for both JWT tokens and API keys. ### Client -Client-side OAuth is implemented by setting -[`StreamableClientTransport.HTTPClient`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk@v0.5.0/mcp#StreamableClientTransport.HTTPClient) to a custom [`http.Client`](https://pkg.go.dev/net/http#Client) -Additional support is forthcoming; see modelcontextprotocol/go-sdk#493. +Client-side authorization is supported via the +[`StreamableClientTransport.OAuthHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#StreamableClientTransport.OAuthHandler) +field. If the handler is provided, the transport will automatically use it to +add an `Authorization: Bearer ` header to every request. The transport +will also call the handler's `Authorize` method if the server returns +`401 Unauthorized` or `403 Forbidden` errors to perform the authorization flow +or facilitate scope step-up authorization. + +The SDK implements the Authorization Code flow in +[`auth.AuthorizationCodeHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#AuthorizationCodeHandler). +This handler supports: + +- [Client ID Metadata Documents](https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#client-id-metadata-documents) +- [Pre-registered clients](https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#preregistration) +- [Dynamic Client Registration](https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#dynamic-client-registration) +- [RFC 9207](https://www.rfc-editor.org/rfc/rfc9207) Authorization Server Issuer Identification + +To use it, configure the handler and assign it to the transport: + +```go +authHandler, _ := auth.NewAuthorizationCodeHandler(&auth.AuthorizationCodeHandlerConfig{ + RedirectURL: "https://myapp.com/oauth2-callback", + // Configure one of the following: + // ClientIDMetadataDocumentConfig: ... + // PreregisteredClientConfig: ... + // DynamicClientRegistrationConfig: ... + AuthorizationCodeFetcher: func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + // Open the args.URL in a browser and return the resulting code, state, and iss. + // See full example in examples/auth/client/main.go. + code := ... + state := ... + iss := ... // "iss" query parameter from the redirect URI (RFC 9207) + return &auth.AuthorizationResult{Code: code, State: state, Iss: iss}, nil + }, +}) + +transport := &mcp.StreamableClientTransport{ + Endpoint: "https://example.com/mcp", + OAuthHandler: authHandler, +} +client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) +session, err := client.Connect(ctx, transport, nil) +``` + +The `auth.AuthorizationCodeHandler` automatically manages token refreshing (if the server provides a refresh token) and step-up authentication (when the server returns `insufficient_scope` error). + +#### Enterprise Managed Authorization (SEP-990) + +For enterprise SSO scenarios where users authenticate with an enterprise Identity Provider (IdP), +the SDK provides +[`extauth.EnterpriseHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth/extauth#EnterpriseHandler), +an implementation of `OAuthHandler` that automates the Enterprise Managed Authorization flow: + +1. **OIDC Login**: User authenticates with enterprise IdP → ID Token +2. **Token Exchange** (RFC 8693): ID Token → ID-JAG at IdP +3. **JWT Bearer Grant** (RFC 7523): ID-JAG → Access Token at MCP Server + +To use enterprise managed authorization, create an `EnterpriseHandler` and assign it to your transport: + +```go +// Create ID token fetcher using OIDC login +idTokenFetcher := func(ctx context.Context) (*oauth2.Token, error) { + oidcConfig := &extauth.OIDCLoginConfig{ + IssuerURL: "https://company.okta.com", + Credentials: &oauthex.ClientCredentials{ + ClientID: "idp-client-id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "idp-client-secret", + }, + }, + RedirectURL: "http://localhost:3142", + Scopes: []string{"openid", "profile", "email"}, + } + + tokens, err := extauth.PerformOIDCLogin(ctx, oidcConfig, authCodeFetcher) + if err != nil { + return nil, err + } + + return tokens, nil +} + +// Create Enterprise Handler +enterpriseHandler, err := extauth.NewEnterpriseHandler(&extauth.EnterpriseHandlerConfig{ + IdPIssuerURL: "https://company.okta.com", + IdPCredentials: &oauthex.ClientCredentials{ + ClientID: "idp-client-id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "idp-client-secret", + }, + }, + MCPAuthServerURL: "https://auth.mcpserver.example", + MCPResourceURI: "https://mcp.mcpserver.example", + MCPCredentials: &oauthex.ClientCredentials{ + ClientID: "mcp-client-id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "mcp-client-secret", + }, + }, + MCPScopes: []string{"read", "write"}, + IDTokenFetcher: idTokenFetcher, +}) + +// Use with transport +transport := &mcp.StreamableClientTransport{ + Endpoint: "https://example.com/mcp", + OAuthHandler: enterpriseHandler, +} +client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) +session, err := client.Connect(ctx, transport, nil) +``` + +The `EnterpriseHandler` automatically manages the token exchange flow. Note that it intentionally does not support refresh tokens - when an access token expires, the entire authorization flow is repeated to ensure enterprise policies are consistently enforced. + +For a complete working example, see [examples/auth/enterprise](https://github.com/modelcontextprotocol/go-sdk/tree/main/examples/auth/enterprise). ## Security Here we discuss the mitigations described under -the MCP spec's [Security Best Practices](https://modelcontextprotocol.io/specification/2025-06-18/basic/security_best_practices) section, and how we handle them. +the MCP's [Security Best Practices](https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices) section, and how we handle them. ### Confused Deputy -The [mitigation](https://modelcontextprotocol.io/specification/2025-06-18/basic/security_best_practices#mitigation), obtaining user consent for dynamically registered clients, -happens on the MCP client. At present we don't provide client-side OAuth support. - +The [mitigation](https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#mitigation), +obtaining user consent for dynamically registered clients, is mostly the +responsibility of the MCP Proxy server implementation. The SDK client does +generate cryptographically secure random `state` values for each authorization +request by default and validates them when the authorization code is returned. +Mismatched state values will result in an error. ### Token Passthrough -The [mitigation](https://modelcontextprotocol.io/specification/2025-06-18/basic/security_best_practices#mitigation-2), accepting only tokens that were issued for the server, depends on the structure +The [mitigation](https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#mitigation-2), accepting only tokens that were issued for the server, depends on the structure of tokens and is the responsibility of the [`TokenVerifier`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#TokenVerifier) provided to [`RequireBearerToken`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#RequireBearerToken). +### Server-Side Request Forgery + +The OAuth discovery helpers apply these [mitigations](https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#mitigation-3) by default: + +- _Enforce HTTPS_. Reject `http://` URLs (except loopback: `localhost`, `127.0.0.1`, `::1`) +and redirects that downgrade `https` to a non-`https` scheme. + +- _Block Private IP Ranges_. Reject targets that resolve to a non-public address +(private, link-local including the `169.254.169.254` metadata endpoint, CGNAT, multicast, +unspecified). The check runs on the initial URL and at dial time, guarding against DNS rebinding. + +- _Validate Redirect Targets_. Reject redirects to non-public addresses and cap the redirect count. + +- _Use Egress Proxies_. Out of scope for the SDK; configure separately. + +**Opting out.** Passing a custom `http.Client` can bypass these checks, making SSRF +protection your responsibility. Specifically: a custom `Transport.DialContext`/`DialTLSContext`, +a non-`*http.Transport`, or a configured `Proxy` disables the dial-time IP check; setting +`CheckRedirect` replaces all redirect validation. In these cases, implement your own IP +blocking, redirect validation, and DNS pinning. + ### Session Hijacking -The [mitigations](https://modelcontextprotocol.io/specification/2025-06-18/basic/security_best_practices#mitigation-3) are as follows: +The [mitigations](https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#mitigation-4) are as follows: - _Verify all inbound requests_. The [`RequireBearerToken`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#RequireBearerToken) middleware function will verify all HTTP requests that it receives. It is the @@ -299,12 +669,36 @@ user's responsibility to wrap that function around all handlers in their server. - _Secure session IDs_. This SDK generates cryptographically secure session IDs by default. If you create your own with [`ServerOptions.GetSessionID`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.GetSessionID), it is your responsibility to ensure they are secure. -If you are using Go 1.24 or above, -we recommend using [`crypto/rand.Text`](https://pkg.go.dev/crypto/rand#Text) +We recommend using [`crypto/rand.Text`](https://pkg.go.dev/crypto/rand#Text). -- _Binding session IDs to user information_. This is an application requirement, out of scope -for the SDK. You can create your own session IDs by setting -[`ServerOptions.GetSessionID`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.GetSessionID). +- _Binding session IDs to user information_. The SDK supports this mitigation through +[`TokenInfo.UserID`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#TokenInfo.UserID). +When a [`TokenVerifier`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#TokenVerifier) +sets `UserID` on the returned `TokenInfo`, the streamable transport will: + 1. Store the user ID when a new session is created. + 2. Verify that subsequent requests to that session include a token with the same `UserID`. + 3. Reject requests with a 403 Forbidden if the user ID doesn't match. + + **Recommendation**: If your `TokenVerifier` can extract a user identifier from the token + (such as a `sub` claim in a JWT, or a user ID associated with an API key), set + `TokenInfo.UserID` to enable this protection. This prevents an attacker with a valid + token from hijacking another user's session by guessing or obtaining their session ID. + +### Issuer Mix-Up + +The [mitigation](https://www.rfc-editor.org/rfc/rfc9207) against issuer mix-up attacks is +implemented per [RFC 9207](https://www.rfc-editor.org/rfc/rfc9207). The SDK client validates +the `iss` parameter in authorization responses to ensure they originated from the expected +authorization server: + +- If `iss` is present in the redirect URI, the SDK verifies it matches the issuer from the + authorization server's metadata. A mismatch results in an error. +- If `iss` is absent but the authorization server advertises + `authorization_response_iss_parameter_supported: true` in its [RFC 8414](https://www.rfc-editor.org/rfc/rfc8414) + metadata, the SDK rejects the response with an error. + +The `AuthorizationCodeFetcher` is responsible for extracting the `iss` query parameter from +the redirect URI and returning it in [`AuthorizationResult.Iss`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#AuthorizationResult). ## Utilities @@ -394,6 +788,13 @@ session if the ping fails, set or [`ServerOptions.KeepAlive`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.KeepAlive). +> **Note**: `ping` is removed from the protocol as of `2026-07-28` by +> [SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575). +> The SDK preserves the API for backward compatibility with legacy clients; +> when both peers negotiate `2026-07-28`, the server returns +> `MethodNotFound` (`-32601`) for `ping` and `KeepAlive` should not be +> enabled. + ### Progress [Progress](https://modelcontextprotocol.io/specification/2025-06-18/basic/utilities/progress) @@ -454,3 +855,23 @@ func Example_progress() { // frobbing widgets 2/2 } ``` + +### Error codes + +The SDK uses the standard JSON-RPC base codes (parse error `-32700`, +invalid request `-32600`, method not found `-32601`, invalid params +`-32602`, internal error `-32603`) plus the following MCP-specific codes: + +| Constant | Code | Meaning | +|---|---|---| +| [`CodeHeaderMismatch`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#CodeHeaderMismatch) | `-32020` | An MCP HTTP header (`Mcp-Method`, `Mcp-Name`, `Mcp-Protocol-Version`, or `Mcp-Param-*`) does not match the JSON-RPC body | +| [`CodeMissingRequiredClientCapabilities`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#CodeMissingRequiredClientCapabilities) | `-32021` | Server requires client capabilities the client did not declare | +| [`CodeUnsupportedProtocolVersion`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#CodeUnsupportedProtocolVersion) | `-32022` | Requested protocol version not supported. Data: `UnsupportedProtocolVersionData{Supported []string, Requested string}` | +| [`CodeResourceNotFound`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#CodeResourceNotFound) | `-32602` | Resource URI not found (variable; was `-32002` before SEP-2164) | + +The error code allocation policy defined in `2026-07-28` +([SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575)) +partitions the JSON-RPC server-error range: + +- `-32000` to `-32019`: implementation-defined; existing SDK usage is grandfathered. +- `-32020` to `-32099`: reserved for the MCP specification. diff --git a/docs/quick_start.md b/docs/quick_start.md new file mode 100644 index 0000000..995b378 --- /dev/null +++ b/docs/quick_start.md @@ -0,0 +1,105 @@ + +# Quick Start + +1. [Installation](#installation) +1. [Getting started](#getting-started) + +## Installation + +``` +go get github.com/modelcontextprotocol/go-sdk/mcp +``` + +## Getting started + +To get started creating an MCP server, create an `mcp.Server` instance, add +features to it, and then run it over an `mcp.Transport`. For example, this +server adds a single simple tool, and then connects it to clients over +stdin/stdout: + +```go +package main + +import ( + "context" + "log" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +type Input struct { + Name string `json:"name" jsonschema:"the name of the person to greet"` +} + +type Output struct { + Greeting string `json:"greeting" jsonschema:"the greeting to tell to the user"` +} + +func SayHi(ctx context.Context, req *mcp.CallToolRequest, input Input) ( + *mcp.CallToolResult, + Output, + error, +) { + return nil, Output{Greeting: "Hi " + input.Name}, nil +} + +func main() { + // Create a server with a single tool. + server := mcp.NewServer(&mcp.Implementation{Name: "greeter", Version: "v1.0.0"}, nil) + mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi) + // Run the server over stdin/stdout, until the client disconnects. + if err := server.Run(context.Background(), &mcp.StdioTransport{}); err != nil { + log.Fatal(err) + } +} +``` + +To communicate with that server, create an `mcp.Client` and connect it to the +corresponding server, by running the server command and communicating over its +stdin/stdout: + +```go +package main + +import ( + "context" + "log" + "os/exec" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +func main() { + ctx := context.Background() + + // Create a new client, with no features. + client := mcp.NewClient(&mcp.Implementation{Name: "mcp-client", Version: "v1.0.0"}, nil) + + // Connect to a server over stdin/stdout. + transport := &mcp.CommandTransport{Command: exec.Command("myserver")} + session, err := client.Connect(ctx, transport, nil) + if err != nil { + log.Fatal(err) + } + defer session.Close() + + // Call a tool on the server. + params := &mcp.CallToolParams{ + Name: "greet", + Arguments: map[string]any{"name": "you"}, + } + res, err := session.CallTool(ctx, params) + if err != nil { + log.Fatalf("CallTool failed: %v", err) + } + if res.IsError { + log.Fatal("tool failed") + } + for _, c := range res.Content { + log.Print(c.(*mcp.TextContent).Text) + } +} +``` + +The [`examples/`](https://github.com/modelcontextprotocol/go-sdk/tree/main/examples) directory contains more example clients and +servers. diff --git a/docs/rough_edges.md b/docs/rough_edges.md new file mode 100644 index 0000000..c1a618d --- /dev/null +++ b/docs/rough_edges.md @@ -0,0 +1,70 @@ + +# Rough edges: API decisions to reconsider for v2 + +This file collects a list of API oversights or rough edges that we've uncovered +post v1.0.0, along with their current workarounds. These issues can't be +addressed without breaking backward compatibility, but we'll revisit them for +v2. + +- `EventStore.Open` is unnecessary. This was an artifact of an earlier version + of the SDK where event persistence and delivery were combined. + + **Workaround**: `Open` may be implemented as a no-op. + +- `Event` need not have been exported: it's an implementation detail of the SSE + and streamable transports. Also the 'Name' field is a misnomer: it should be + 'event'. + +- Enforcing valid tool names: with + [SEP-986](https://github.com/modelcontextprotocol/modelcontextprotocol/issues/986) + landing after the SDK was at v1, we missed an opportunity to panic on invalid + tool names. Instead, we have to simply produce an error log. In v2, we should + panic. + +- Inconsistent naming. + - `ResourceUpdatedNotificationsParams` should probably have just been + `ResourceUpdatedParams`, as we don't include the word 'notification' in + other notification param types. + - Similarly, `ProgressNotificationParams` should probably have been + `ProgressParams`. + +- `AudioContent.MarshalJSON` should have had a pointer receiver, to be + consistent with other content types. + +- `ClientCapabilities.Roots` should have been a distinguished struct pointer + ([see #607](https://github.com/modelcontextprotocol/go-sdk/issues/607)). + + **Workaround**: use `ClientCapabilities.RootsV2`, which aligns with the + semantics of other capability fields. + +- Default capabilities should have been empty. Instead, servers default to + advertising `logging`, and clients default to advertising `roots` with + `listChanged: true`. This is confusing because a nil `Capabilities` field + does not mean "no capabilities". + + **Workaround**: to advertise no capabilities, set + `ServerOptions.Capabilities` or `ClientOptions.Capabilities` to an empty + `&ServerCapabilities{}` or `&ClientCapabilities{}` respectively. + +- `CreateMessageResult.Content` is singular `Content`, but the 2025-11-25 spec + allows `content` to be a single block or an array (for parallel tool calls). + We added `CreateMessageResultWithTools` (with `Content []Content`) as a + workaround, matching the TypeScript SDK's approach. In v2, + `CreateMessageResult` should use `[]Content` directly. Similarly, + `SamplingMessage.Content` should become `[]Content` to support sending + multiple tool_result blocks in a single user message. + +- We didn't actually need CallToolParams and CallToolParamsRaw, since even when + we're unmarshalling into a custom Go type (for the mcp.AddTool convenience + wrapper) we need to first unmarshal into a `map[string]any` in order to do + server-side validation of required fields. CallToolParams could have just had + a map[string]any. + +- `StreamableHTTPOptions.CrossOriginProtection` should not have been part of + the SDK API. Cross-origin protection is a general HTTP concern, not specific + to MCP, and can be applied as standard HTTP middleware. + +- `ToolAnnotations` (`mcp/protocol.go`) should have all fields typed as `*bool` + for full control to define what is being sent over the wire. Different + MCP clients have different requirements, and some of them require all fields + to be explicitly set to either `true` or `false`. diff --git a/docs/server.md b/docs/server.md index f986c68..8973e3a 100644 --- a/docs/server.md +++ b/docs/server.md @@ -2,11 +2,24 @@ # Support for MCP server features 1. [Prompts](#prompts) + 1. [Prompt message content](#prompt-message-content) 1. [Resources](#resources) + 1. [Binary resources](#binary-resources) + 1. [Resource subscriptions](#resource-subscriptions) 1. [Tools](#tools) + 1. [Tool result content](#tool-result-content) +1. [List changed notifications](#list-changed-notifications) +1. [Multi Round-Trip Requests](#multi-round-trip-requests) + 1. [Talking to legacy clients](#talking-to-legacy-clients) + 1. [Example](#example) +1. [Cacheable list results](#cacheable-list-results) 1. [Utilities](#utilities) 1. [Completion](#completion) 1. [Logging](#logging) +1. [Capabilities](#capabilities) + 1. [Capability inference](#capability-inference) + 1. [Explicit capabilities](#explicit-capabilities) + 1. [Extensions](#extensions) 1. [Pagination](#pagination) ## Prompts @@ -107,6 +120,91 @@ func Example_prompts() { } ``` +### Prompt message content + +A +[`PromptMessage`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#PromptMessage) +has two fields: a `Role`, either `"user"` or `"assistant"`, and a single +`Content`. That is the same interface a tool result uses (see [Tool result +content](#tool-result-content)), so a prompt message can hold an image, audio, +or a resource's contents as readily as text. + +Media on its own says nothing about what the model should do with it, so send a +text message alongside it. Embedding a resource puts its contents directly in +the message, sparing the client a separate `resources/read`. Note that +`ResourceContents.URI` is not checked against the resources the server has +registered: it records where the contents came from, so use the URI of a real +resource when there is one, as in the example below. + +```go +func Example_promptContent() { + ctx := context.Background() + + const styleGuide = "- Prefer clarity over cleverness.\n" + screenshotPNG := []byte("\x89PNG\r\n\x1a\n") + + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + + s.AddResource(&mcp.Resource{URI: "doc://style-guide", MIMEType: "text/markdown"}, + func(context.Context, *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{ + Contents: []*mcp.ResourceContents{{ + URI: "doc://style-guide", + MIMEType: "text/markdown", + Text: styleGuide, + }}, + }, nil + }) + + s.AddPrompt(&mcp.Prompt{Name: "review_screenshot"}, + func(context.Context, *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{ + Messages: []*mcp.PromptMessage{ + {Role: "user", Content: &mcp.EmbeddedResource{ + Resource: &mcp.ResourceContents{ + URI: "doc://style-guide", + MIMEType: "text/markdown", + Text: styleGuide, + }, + }}, + {Role: "user", Content: &mcp.ImageContent{Data: screenshotPNG, MIMEType: "image/png"}}, + {Role: "user", Content: &mcp.TextContent{Text: "Review the screenshot against the style guide."}}, + }, + }, nil + }) + + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) + t1, t2 := mcp.NewInMemoryTransports() + if _, err := s.Connect(ctx, t1, nil); err != nil { + log.Fatal(err) + } + cs, err := c.Connect(ctx, t2, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + res, err := cs.GetPrompt(ctx, &mcp.GetPromptParams{Name: "review_screenshot"}) + if err != nil { + log.Fatal(err) + } + for _, msg := range res.Messages { + switch content := msg.Content.(type) { + case *mcp.TextContent: + fmt.Println(msg.Role, "text:", content.Text) + case *mcp.ImageContent: + fmt.Printf("%s image: %s, %d bytes\n", msg.Role, content.MIMEType, len(content.Data)) + case *mcp.EmbeddedResource: + fmt.Printf("%s embedded resource: %s (%s)\n", msg.Role, content.Resource.URI, content.Resource.MIMEType) + } + } + // Output: + // user embedded resource: doc://style-guide (text/markdown) + // user image: image/png, 8 bytes + // user text: Review the screenshot against the style guide. +} +``` + ## Resources In MCP terms, a _resource_ is some data referenced by a URI. @@ -139,7 +237,12 @@ and to unsubscribe. Set [`ClientOptions.ResourceUpdatedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.ResourceUpdatedHandler) -to be notified of changes to subscribed resources. +to be notified of changes to subscribed resources. On `2026-07-28` and +later sessions the SDK delivers these notifications over a +`subscriptions/listen` stream instead of the legacy `resources/subscribe` +RPC; see [Subscriptions +(`subscriptions/listen`)](protocol.md#subscriptions) +for the wire-level details. **Server-side**: Use @@ -233,6 +336,145 @@ func Example_resources() { } ``` +### Binary resources + +A +[`ResourceContents`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ResourceContents) +carries its data in either `Text` or `Blob`. Binary data goes in `Blob` as +plain bytes — the SDK base64-encodes it on the wire — and leaves `Text` empty, +so a client tells the two apart by which field is populated. + +```go +func Example_binaryResource() { + ctx := context.Background() + + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + s.AddResource(&mcp.Resource{URI: "file:///logo.png", MIMEType: "image/png"}, + func(context.Context, *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{ + Contents: []*mcp.ResourceContents{{ + URI: "file:///logo.png", + MIMEType: "image/png", + Blob: []byte("\x89PNG\r\n\x1a\n"), + }}, + }, nil + }) + + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) + t1, t2 := mcp.NewInMemoryTransports() + if _, err := s.Connect(ctx, t1, nil); err != nil { + log.Fatal(err) + } + cs, err := c.Connect(ctx, t2, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + res, err := cs.ReadResource(ctx, &mcp.ReadResourceParams{URI: "file:///logo.png"}) + if err != nil { + log.Fatal(err) + } + contents := res.Contents[0] + fmt.Printf("%s: %d bytes of %s, text is %q\n", + contents.URI, len(contents.Blob), contents.MIMEType, contents.Text) + + // Output: + // file:///logo.png: 8 bytes of image/png, text is "" +} +``` + +### Resource subscriptions + +A client interested in changes to one resource subscribes to its URI with +[`ClientSession.Subscribe`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientSession.Subscribe) +and stops with +[`ClientSession.Unsubscribe`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientSession.Unsubscribe). +Set +[`ServerOptions.SubscribeHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.SubscribeHandler) +and +[`UnsubscribeHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.UnsubscribeHandler) +to track who is listening; setting either one gives the server the +`resources.subscribe` capability. + +Publish a change with +[`Server.ResourceUpdated`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server.ResourceUpdated), +which notifies every session subscribed to that URI. The notification carries +only the URI, not the new contents, so a client that wants them reads the +resource again. + +```go +func Example_resourceSubscription() { + ctx := context.Background() + + config := "theme=light\n" + + updated := make(chan string, 1) + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, &mcp.ClientOptions{ + // The notification carries only the URI, so read the resource again to + // see what it now holds. + ResourceUpdatedHandler: func(context.Context, *mcp.ResourceUpdatedNotificationRequest) { + updated <- "config://app" + }, + }) + + subscribed := make(chan string, 2) + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, &mcp.ServerOptions{ + SubscribeHandler: func(_ context.Context, req *mcp.SubscribeRequest) error { + subscribed <- "subscribed to " + req.Params.URI + return nil + }, + UnsubscribeHandler: func(_ context.Context, req *mcp.UnsubscribeRequest) error { + subscribed <- "unsubscribed from " + req.Params.URI + return nil + }, + }) + s.AddResource(&mcp.Resource{URI: "config://app"}, + func(context.Context, *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{ + Contents: []*mcp.ResourceContents{{URI: "config://app", Text: config}}, + }, nil + }) + + t1, t2 := mcp.NewInMemoryTransports() + if _, err := s.Connect(ctx, t1, nil); err != nil { + log.Fatal(err) + } + cs, err := c.Connect(ctx, t2, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + if err := cs.Subscribe(ctx, &mcp.SubscribeParams{URI: "config://app"}); err != nil { + log.Fatal(err) + } + fmt.Println(<-subscribed) + + config = "theme=dark\n" + if err := s.ResourceUpdated(ctx, &mcp.ResourceUpdatedNotificationParams{URI: "config://app"}); err != nil { + log.Fatal(err) + } + uri := <-updated + + res, err := cs.ReadResource(ctx, &mcp.ReadResourceParams{URI: uri}) + if err != nil { + log.Fatal(err) + } + fmt.Print("re-read: ", res.Contents[0].Text) + + if err := cs.Unsubscribe(ctx, &mcp.UnsubscribeParams{URI: "config://app"}); err != nil { + log.Fatal(err) + } + fmt.Println(<-subscribed) + + // Output: + // subscribed to config://app + // re-read: theme=dark + // unsubscribed from config://app +} +``` + ## Tools MCP servers can provide @@ -284,7 +526,7 @@ In order to implement a tool, the user must do all of the following: - Unmarshal the input schema into a Go value - Execute the tool logic. - Marshal the tool's structured output (if any) to JSON, and store it in the - result's `StructuredOutput` field as well as the unstructured `Content` field. + result's `StructuredContent` field as well as the unstructured `Content` field. - Validate that output JSON against the tool's output schema. - If any tool errors occurred, pack them into the unstructured content and set `IsError` to `true.` @@ -309,13 +551,17 @@ mcp.AddTool(server, &mcp.Tool{Name: "my_tool"}, handler) This does the following automatically: -- If `Tool.InputSchema` or `Tool.OutputSchema` are unset, the input and output - schemas are inferred from the `In` type, which must be a struct or map. - Optional `jsonschema` struct tags provide argument descriptions. +- If `Tool.InputSchema` is unset, the input schema is inferred from the `In` + type, which must be a struct or map. +- If `Tool.OutputSchema` is unset and the `Out` type is not `any`, the output + schema is inferred from the `Out` type. Per SEP-2106, `Out` may be any Go + type whose inferred schema is a valid JSON Schema (struct, map, slice, + primitive, etc.). +- Optional `jsonschema` struct tags provide argument and output descriptions. - Tool arguments are validated against the input schema. - Tool arguments are marshaled into the `In` value. - Tool output (the `Out` value) is marshaled into the result's - `StructuredOutput`, as well as the unstructured `Content`. + `StructuredContent`, as well as the unstructured `Content`. - Output is validated against the tool's output schema. - If an ordinary error is returned, it is stored int the `CallToolResult` and `IsError` is set to `true`. @@ -396,10 +642,366 @@ mcp.AddTool(server, &mcp.Tool{ }, WeatherTool) ``` -_See [mcp/tool_example_test.go](../mcp/tool_example_test.go) for the full -example, or [examples/server/toolschemas](examples/server/toolschemas/main.go) +_See [mcp/tool_example_test.go](https://github.com/modelcontextprotocol/go-sdk/blob/main/mcp/tool_example_test.go) for the full +example, or [examples/server/toolschemas](https://github.com/modelcontextprotocol/go-sdk/blob/main/examples/server/toolschemas/main.go) for more examples of customizing tool schemas._ +### Tool result content + +Alongside its structured output, a tool result carries a list of +[`Content`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Content) +blocks in `CallToolResult.Content`, and may mix as many kinds as it needs: + +| Type | Carries | +| --- | --- | +| [`TextContent`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#TextContent) | plain text | +| [`ImageContent`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ImageContent) | image data, with a MIME type | +| [`AudioContent`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#AudioContent) | audio data, with a MIME type | +| [`EmbeddedResource`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#EmbeddedResource) | the contents of a resource, inline | +| [`ResourceLink`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ResourceLink) | a reference to a resource the client can read separately | + +`ImageContent.Data` and `AudioContent.Data` are plain `[]byte`; the SDK +base64-encodes them on the wire, so handlers never encode by hand, and base64 +carried over from another MCP SDK must be decoded before it is assigned. An +`EmbeddedResource` wraps a +[`ResourceContents`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ResourceContents), +which holds either `Text` or, for binary data, `Blob`. Prefer a `ResourceLink` +when the client may not need the contents, since an embedded resource is +transferred whether it is used or not. + +When a handler bound with `AddTool` leaves `Content` unset, the SDK fills it +with the JSON encoding of the output value; set `Content` explicitly, as below, +to return anything else. `StructuredContent` is still populated from the output +value either way. + +```go +func ExampleAddTool_contentTypes() { + chartPNG := []byte("\x89PNG\r\n\x1a\n") + summaryWAV := []byte("RIFF....WAVE") + + server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(server, &mcp.Tool{Name: "report"}, func(context.Context, *mcp.CallToolRequest, struct{}) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "Q3 revenue by region."}, + &mcp.ImageContent{Data: chartPNG, MIMEType: "image/png"}, + &mcp.AudioContent{Data: summaryWAV, MIMEType: "audio/wav"}, + // An embedded resource inlines the contents, so the client + // needs no follow-up resources/read. Binary contents go in + // ResourceContents.Blob instead of Text. + &mcp.EmbeddedResource{ + Resource: &mcp.ResourceContents{ + URI: "file:///reports/q3.csv", + MIMEType: "text/csv", + Text: "region,revenue\nEMEA,42\n", + }, + }, + // A resource link only points at a resource. Prefer it when the + // client may not need the contents, since an embedded resource + // is transferred either way. + &mcp.ResourceLink{URI: "file:///reports/q3.pdf", MIMEType: "application/pdf"}, + }, + }, nil, nil + }) + + ctx := context.Background() + session, err := connect(ctx, server) + if err != nil { + log.Fatal(err) + } + defer session.Close() + + res, err := session.CallTool(ctx, &mcp.CallToolParams{Name: "report"}) + if err != nil { + log.Fatal(err) + } + + // Clients type-switch over the content they receive. + for _, content := range res.Content { + switch c := content.(type) { + case *mcp.TextContent: + fmt.Println("text:", c.Text) + case *mcp.ImageContent: + fmt.Printf("image: %s, %d bytes\n", c.MIMEType, len(c.Data)) + case *mcp.AudioContent: + fmt.Printf("audio: %s, %d bytes\n", c.MIMEType, len(c.Data)) + case *mcp.EmbeddedResource: + fmt.Printf("embedded resource: %s (%s)\n", c.Resource.URI, c.Resource.MIMEType) + case *mcp.ResourceLink: + fmt.Printf("resource link: %s (%s)\n", c.URI, c.MIMEType) + } + } + // Output: + // text: Q3 revenue by region. + // image: image/png, 8 bytes + // audio: audio/wav, 12 bytes + // embedded resource: file:///reports/q3.csv (text/csv) + // resource link: file:///reports/q3.pdf (application/pdf) +} +``` + +`ToolUseContent` and `ToolResultContent` also implement `Content`, but are only +valid in sampling messages, not in tool results. + +**Stateless server deployments:** Some deployments create a new +[`Server`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server) +for each incoming request, re-registering tools every time. To avoid repeated +schema generation, create a +[`SchemaCache`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#SchemaCache) +and share it across server instances: + +```go +var schemaCache = mcp.NewSchemaCache() // create once at startup + +func handleRequest(w http.ResponseWriter, r *http.Request) { + s := mcp.NewServer(impl, &mcp.ServerOptions{SchemaCache: schemaCache}) + mcp.AddTool(s, myTool, myHandler) + // ... +} +``` + +## List changed notifications + +Adding or removing a feature on a connected server sends the matching +`notifications/*/list_changed` to every client, which dispatches it to +[`ClientOptions.ToolListChangedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.ToolListChangedHandler), +[`PromptListChangedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.PromptListChangedHandler), +or +[`ResourceListChangedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.ResourceListChangedHandler). +A notification reports only that the list changed, so a client that needs the +new contents lists them again. + +These notifications ride on capabilities, and capabilities are inferred from +what is registered before [`Server.Connect`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server.Connect) +(see [Capabilities](#capabilities)). A server that registers its features +afterwards, as the example below does, has to declare +`ServerOptions.HasTools`, `HasPrompts`, or `HasResources` itself, or the +client is never told. + +```go +func Example_listChanged() { + ctx := context.Background() + + changed := make(chan string, 3) + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, &mcp.ClientOptions{ + ToolListChangedHandler: func(context.Context, *mcp.ToolListChangedRequest) { + changed <- "tools" + }, + PromptListChangedHandler: func(context.Context, *mcp.PromptListChangedRequest) { + changed <- "prompts" + }, + ResourceListChangedHandler: func(context.Context, *mcp.ResourceListChangedRequest) { + changed <- "resources" + }, + }) + + // Nothing is registered before Connect, so the capabilities that carry these + // notifications have to be declared explicitly. + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, &mcp.ServerOptions{ + HasTools: true, + HasPrompts: true, + HasResources: true, + }) + + t1, t2 := mcp.NewInMemoryTransports() + if _, err := s.Connect(ctx, t1, nil); err != nil { + log.Fatal(err) + } + cs, err := c.Connect(ctx, t2, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + mcp.AddTool(s, &mcp.Tool{Name: "greet"}, func(context.Context, *mcp.CallToolRequest, struct{}) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{}, nil, nil + }) + fmt.Println("changed:", <-changed) + + s.AddPrompt(&mcp.Prompt{Name: "review"}, func(context.Context, *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{}, nil + }) + fmt.Println("changed:", <-changed) + + s.AddResource(&mcp.Resource{URI: "file:///a"}, func(context.Context, *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{}, nil + }) + fmt.Println("changed:", <-changed) + + // Removing a feature notifies the client just as adding one does. + s.RemoveTools("greet") + fmt.Println("changed:", <-changed) + + // Output: + // changed: tools + // changed: prompts + // changed: resources + // changed: tools +} +``` + +## Multi Round-Trip Requests + +[SEP-2322](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2322) +defines a new pattern for server-to-client requests (sampling, elicitation, +roots). Instead of issuing a fresh JSON-RPC request mid-flight, the server +returns an `InputRequiredResult` from its in-flight handler — the +`InputRequests` field of `CallToolResult`, `GetPromptResult`, or +`ReadResourceResult` carries the requests for additional information. The +client responds by retrying the original call with `InputResponses` set. + +The SDK supports this pattern from both sides without requiring callers to +choose: + +- **Returning input requests from a handler.** Set `InputRequests` on the + result you return. + Leave `Content` / `StructuredContent` empty — returning + both at once is a server bug and the SDK returns `-32603 InternalError`. +- **Calling the legacy APIs.** `ServerSession.Elicit`, + `ServerSession.CreateMessage(WithTools)`, and `ServerSession.ListRoots` + remain available and work for both new and legacy clients. + +### Talking to legacy clients + +The server installs `serverMultiRoundTripMiddleware` automatically. For +clients on a protocol version earlier than `2026-07-28`, the middleware +intercepts any `InputRequiredResult` your handler returns, fulfils each +input request itself by calling the legacy server-initiated APIs +(`Elicit`, `CreateMessage`, `ListRoots`), and re-invokes your handler +exactly once with the responses already populated. This means a handler +written in the MRTR style works against both old and new clients without +code changes. + +### Example + +The following example shows a "greet" tool whose handler asks the user +for their name via elicitation before producing the final greeting. The +handler runs twice — once to issue the elicitation, once to consume the +response — but the client call site sees a single `CallTool` returning the +final result, because the SDK's MRTR middleware handles the round trip on +either side of the wire (depending on the negotiated protocol version). + +```go +// Example_mrtr demonstrates the [Multi Round-Trip Requests] pattern +// (SEP-2322). A tool handler signals "I need more information from the user" +// by returning a [mcp.CallToolResult] whose `InputRequests` field carries the +// requests for the additional information. Each request can be an +// [mcp.ElicitParams] (ask the user a question), +// [mcp.CreateMessageParams] (sample from the client's LLM), or +// [mcp.ListRootsParams] (list the client's roots). +// +// On protocol version `2026-07-28` and later, the client's +// `clientMultiRoundTripMiddleware` fulfils each input request from the +// configured handler and retries the original call transparently. On earlier +// protocol versions, the server's `serverMultiRoundTripMiddleware` performs +// the equivalent dance from the server side by calling the legacy +// server-to-client API (`ServerSession.Elicit`, `CreateMessage`, or +// `ListRoots`) and re-invoking the handler with the response. Either way, +// the user-facing client call site sees only the final result. +// +// [Multi Round-Trip Requests]: https://modelcontextprotocol.io/specification/draft/basic/patterns#multi-round-trip-requests +func Example_mrtr() { + ctx := context.Background() + + // Server: a "greet" tool that asks the user for their name before + // returning a greeting. The handler runs twice per logical call: + // once to issue the elicitation, once to consume the response. + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(s, &mcp.Tool{Name: "greet"}, func(_ context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + // First call: ask the user for their name. The map key + // ("user_name") is a label the handler chooses; the client must + // echo it back in InputResponses. The wire-level method name + // ("elicitation/create") is derived from the value's Go type. + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{ + "user_name": &mcp.ElicitParams{ + Message: "What's your name?", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "name": {Type: "string"}, + }, + }, + }, + }, + // RequestState is an opaque token the client echoes back on + // the retry, letting the handler resume its work without + // per-session storage. + RequestState: "step=1", + }, nil, nil + } + // Retry: read the elicitation response and produce the final greeting. + name := req.Params.InputResponses["user_name"].(*mcp.ElicitResult).Content["name"].(string) + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "Hello " + name}}, + }, nil, nil + }) + + // Client: declares an elicitation handler. The SDK's MRTR middleware + // uses it to fulfil any input request the tool handler returns. + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, &mcp.ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + return &mcp.ElicitResult{Action: "accept", Content: map[string]any{"name": "MCP Go"}}, nil + }, + }) + + ct, st := mcp.NewInMemoryTransports() + if _, err := s.Connect(ctx, st, nil); err != nil { + log.Fatal(err) + } + cs, err := c.Connect(ctx, ct, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + // Single call site. The MRTR round trip is invisible to the caller. + res, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: "greet"}) + if err != nil { + log.Fatal(err) + } + fmt.Println(res.Content[0].(*mcp.TextContent).Text) +} +``` + +## Cacheable list results + +[SEP-2549](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2549) +adds `ttlMs` and `cacheScope` fields to the results of `tools/list`, +`prompts/list`, `resources/list`, `resources/templates/list`, and +`resources/read`. Both fields complement existing list +results: clients use `ttlMs` as a freshness hint to reduce polling, +and `cacheScope` (`"public"` or `"private"`) controls whether shared +intermediaries may cache the response. + +**Server-side defaults.** The SDK fills in `CacheScope = "public"` whenever a +result leaves it empty, since the field is required on the wire. Everything +else is left as produced: a `ResourceHandler` that sets `TTLMs` or `CacheScope` +on its `ReadResourceResult` keeps those values. The list results and +`server/discover` are built by the SDK itself, so they start out with +`TTLMs = 0`, which a client cache treats as "immediately stale". + +**Setting a policy.** `ServerOptions.SetCacheable` decides the fields for every +result that carries them. It runs after the corresponding handler, receives +what that handler produced, and may keep, adjust, or replace it. The request is +passed in so that policy can vary by method, session, or authorization context. + +```go +server := mcp.NewServer(impl, &mcp.ServerOptions{ + SetCacheable: func(_ context.Context, req mcp.Request, c *mcp.Cacheable) { + // A 30s freshness hint, except where the handler chose its own. + if c.TTLMs == 0 { + c.TTLMs = 30_000 + } + }, +}) +``` + +`SetCacheable` may run with the server's lock held, so it must not call back +into the `Server`: adding or removing a feature, or ranging over +`Server.Sessions`, deadlocks. + ## Utilities ### Completion @@ -450,6 +1052,13 @@ _ = mcp.NewServer(&mcp.Implementation{Name: "server"}, &mcp.ServerOptions{ ### Logging +> **Note:** The logging feature is deprecated as of protocol version +> 2026-07-28 +> ([SEP-2577](https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging)). +> It remains fully functional during the deprecation window (at least twelve +> months). The SDK continues to support logging for compatibility. Servers +> should migrate to stderr logging (for STDIO transports) or OpenTelemetry. + MCP servers can send logging messages to MCP clients. (This form of logging is distinct from server-side logging, where the server produces logs that remain server-side, for use by server maintainers.) @@ -457,8 +1066,12 @@ server produces logs that remain server-side, for use by server maintainers.) **Server-side**: The minimum log level is part of the server state. For stateful sessions, there is no default log level: no log messages will be sent -until the client calls `SetLevel` (see below). -For stateful sessions, the level defaults to "info". +until the client calls `SetLoggingLevel` (see below). +For legacy stateless sessions the level defaults to "info". +For sessionless 2026-07-28 requests there is no server-side default: each request +carries its level in `_meta` (`io.modelcontextprotocol/logLevel`), and nothing is +logged until the client sends one. +(Note: logging is deprecated as of 2026-07-28 / SEP-2577.) [`ServerSession.Log`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession.Log) is the low-level way for servers to log to clients. It sends a logging notification to the client if the level of the message @@ -474,7 +1087,7 @@ Servers always report the logging capability. **Client-side**: Set [`ClientOptions.LoggingMessageHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.LoggingMessageHandler) to receive log messages. -Call [`ClientSession.SetLevel`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientSession.SetLevel) to change the log level for a session. +Call [`ClientSession.SetLoggingLevel`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientSession.SetLoggingLevel) to change the log level for a session. ```go func Example_logging() { @@ -535,6 +1148,61 @@ func Example_logging() { } ``` +## Capabilities + +Server capabilities are advertised to clients during the initialization +handshake. By default, the SDK advertises only the `logging` capability. +Additional capabilities are automatically added when features are registered +(e.g., adding a tool adds the `tools` capability). + +### Capability inference + +When features such as tools, prompts, or resources are added to the server +(e.g., via `Server.AddTool`), their capability is automatically inferred, with +default value `{listChanged:true}`. Similarly, if the +`ServerOptions.SubscribeHandler` or `ServerOptions.CompletionHandler` are set, +the corresponding capability is added. + +### Explicit capabilities + +To explicitly declare capabilities, or to override the [default inferred +capability](#capability-inference), set +[`ServerOptions.Capabilities`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.Capabilities). +This sets the default server capabilities, before any capabilities are added +based on configured handlers. If a capability is already present as a field in +`Capabilities`, adding a feature or handler will not change its configuration. + +This allows you to: + +- **Disable default capabilities**: Pass an empty `&ServerCapabilities{}` to + disable all defaults, including logging. +- **Disable listChanged notifications**: Set `ListChanged: false` on a + capability to prevent the server from sending list-changed notifications + when features are added or removed. +- **Pre-declare capabilities**: Declare capabilities before features are + registered, useful for servers that load features dynamically. + +```go +// Disable listChanged notifications for tools +server := mcp.NewServer(impl, &mcp.ServerOptions{ + Capabilities: &mcp.ServerCapabilities{ + Logging: &mcp.LoggingCapabilities{}, + Tools: &mcp.ToolCapabilities{ListChanged: false}, + }, +}) +``` + +**Deprecated**: The `HasPrompts`, `HasResources`, and `HasTools` fields on +`ServerOptions` are deprecated. Use `Capabilities` instead. + +### Extensions + +[SEP-2133](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2133) +adds an `extensions` map to `ServerCapabilities` so that optional +capabilities outside the core protocol can be declared on the wire. Keys +are namespaced as `"{vendor-prefix}/{extension-name}"`; values are +per-extension settings objects. + ### Pagination Server-side feature lists may be diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 38410ad..950eef7 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -33,7 +33,7 @@ func ExampleLoggingTransport() { if err != nil { log.Fatal(err) } - defer serverSession.Wait() + defer serverSession.Close() client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) var b bytes.Buffer @@ -50,9 +50,8 @@ func ExampleLoggingTransport() { } // Output: - // read: {"jsonrpc":"2.0","id":1,"result":{"capabilities":{"logging":{}},"protocolVersion":"2025-06-18","serverInfo":{"name":"server","version":"v0.0.1"}}} - // write: {"jsonrpc":"2.0","id":1,"method":"initialize","params":{"capabilities":{"roots":{"listChanged":true}},"clientInfo":{"name":"client","version":"v0.0.1"},"protocolVersion":"2025-06-18"}} - // write: {"jsonrpc":"2.0","method":"notifications/initialized","params":{}} + // read: {"jsonrpc":"2.0","id":1,"result":{"resultType":"complete","_meta":{"io.modelcontextprotocol/serverInfo":{"name":"server","version":"v0.0.1"}},"ttlMs":0,"cacheScope":"public","supportedVersions":["2026-07-28","2025-11-25","2025-06-18","2025-03-26","2024-11-05"],"capabilities":{"logging":{}}}} + // write: {"jsonrpc":"2.0","id":1,"method":"server/discover","params":{"_meta":{"io.modelcontextprotocol/clientCapabilities":{"roots":{"listChanged":true}},"io.modelcontextprotocol/clientInfo":{"name":"client","version":"v0.0.1"},"io.modelcontextprotocol/protocolVersion":"2026-07-28"}}} } ``` @@ -71,7 +70,7 @@ func ExampleStreamableHTTPHandler_middleware() { server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.1.0"}, nil) handler := mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { return server - }, nil) + }, &mcp.StreamableHTTPOptions{Stateless: true}) loggingHandler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { // Example debugging; you could also capture the response. body, err := io.ReadAll(req.Body) diff --git a/examples/auth/client/main.go b/examples/auth/client/main.go new file mode 100644 index 0000000..7e5d464 --- /dev/null +++ b/examples/auth/client/main.go @@ -0,0 +1,132 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package main + +import ( + "context" + "errors" + "flag" + "fmt" + "log" + "net" + "net/http" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +var ( + // URL of the MCP server. + serverURL = flag.String("server_url", "http://localhost:8000/mcp", "URL of the MCP server.") + // Port for the local HTTP server that will receive the authorization code. + callbackPort = flag.Int("callback_port", 3142, "Port for the local HTTP server that will receive the authorization code.") +) + +type codeReceiver struct { + authChan chan *auth.AuthorizationResult + errChan chan error + server *http.Server +} + +func (r *codeReceiver) serveRedirectHandler(listener net.Listener) { + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) { + r.authChan <- &auth.AuthorizationResult{ + Code: req.URL.Query().Get("code"), + State: req.URL.Query().Get("state"), + Iss: req.URL.Query().Get("iss"), + } + fmt.Fprint(w, "Authentication successful. You can close this window.") + }) + + r.server = &http.Server{ + Addr: fmt.Sprintf("localhost:%d", *callbackPort), + Handler: mux, + } + if err := r.server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) { + r.errChan <- err + } +} + +func (r *codeReceiver) getAuthorizationCode(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + fmt.Printf("Please open the following URL in your browser: %s\n", args.URL) + select { + case authRes := <-r.authChan: + return authRes, nil + case err := <-r.errChan: + return nil, err + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (r *codeReceiver) close() { + if r.server != nil { + r.server.Close() + } +} + +func main() { + flag.Parse() + receiver := &codeReceiver{ + authChan: make(chan *auth.AuthorizationResult), + errChan: make(chan error), + } + listener, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", *callbackPort)) + if err != nil { + log.Fatalf("failed to listen: %v", err) + } + go receiver.serveRedirectHandler(listener) + defer receiver.close() + + authHandler, err := auth.NewAuthorizationCodeHandler(&auth.AuthorizationCodeHandlerConfig{ + RedirectURL: fmt.Sprintf("http://localhost:%d", *callbackPort), + AuthorizationCodeFetcher: receiver.getAuthorizationCode, + // Uncomment the client configuration you want to use. + // PreregisteredClient: &oauthex.ClientCredentials{ + // ClientID: "", + // ClientSecretAuth: &oauthex.ClientSecretAuth{ + // ClientSecret: "", + // }, + // }, + // }, + // DynamicClientRegistrationConfig: &auth.DynamicClientRegistrationConfig{ + // Metadata: &oauthex.ClientRegistrationMetadata{ + // ClientName: "Dynamically registered MCP client", + // RedirectURIs: []string{fmt.Sprintf("http://localhost:%d", *callbackPort)}, + // Scope: "read", + // }, + // }, + }) + if err != nil { + log.Fatalf("failed to create auth handler: %v", err) + } + + transport := &mcp.StreamableClientTransport{ + Endpoint: *serverURL, + OAuthHandler: authHandler, + } + + ctx := context.Background() + client := mcp.NewClient(&mcp.Implementation{ + Name: "test-client", + Version: "1.0.0", + }, nil) + + session, err := client.Connect(ctx, transport, nil) + if err != nil { + log.Fatalf("client.Connect(): %v", err) + } + defer session.Close() + + tools, err := session.ListTools(ctx, nil) + if err != nil { + log.Fatalf("session.ListTools(): %v", err) + } + log.Println("Tools:") + for _, tool := range tools.Tools { + log.Printf("- %q", tool.Name) + } +} diff --git a/examples/auth/enterprise/main.go b/examples/auth/enterprise/main.go new file mode 100644 index 0000000..fb58aad --- /dev/null +++ b/examples/auth/enterprise/main.go @@ -0,0 +1,226 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package main + +import ( + "context" + "errors" + "flag" + "fmt" + "log" + "net" + "net/http" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/auth/extauth" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" +) + +var ( + // IdP (Identity Provider) configuration. + idpIssuerURL = flag.String("idp_issuer", "https://your-idp.okta.com", "IdP issuer URL (e.g., https://your-company.okta.com)") + idpClientID = flag.String("idp_client_id", "", "Client ID registered at the IdP") + idpClientSecret = flag.String("idp_client_secret", "", "Client secret at the IdP (optional for public clients)") + + // MCP Server configuration. + mcpServerURL = flag.String("mcp_server", "http://localhost:8000/mcp", "URL of the MCP server") + mcpAuthServerURL = flag.String("mcp_auth_server", "https://auth.mcpserver.example", "MCP server's authorization server URL") + mcpResourceURI = flag.String("mcp_resource_uri", "https://mcp.mcpserver.example", "MCP server's resource identifier (RFC 9728)") + mcpClientID = flag.String("mcp_client_id", "", "Client ID at the MCP server (optional)") + mcpClientSecret = flag.String("mcp_client_secret", "", "Client secret at the MCP server (optional)") + + // OAuth callback configuration. + callbackPort = flag.Int("callback_port", 3142, "Port for the local HTTP server that will receive the OAuth callback") +) + +// codeReceiver handles the OAuth callback from the IdP's authorization endpoint. +// It starts a local HTTP server to receive the authorization code after the user +// authenticates with their enterprise IdP. +type codeReceiver struct { + authChan chan *auth.AuthorizationResult + errChan chan error + server *http.Server +} + +// serveRedirectHandler starts an HTTP server to handle the OAuth redirect callback. +func (r *codeReceiver) serveRedirectHandler(listener net.Listener) { + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) { + // Extract the authorization code and state from the callback URL. + r.authChan <- &auth.AuthorizationResult{ + Code: req.URL.Query().Get("code"), + State: req.URL.Query().Get("state"), + } + fmt.Fprint(w, "Authentication successful. You can close this window.") + }) + + r.server = &http.Server{ + Addr: fmt.Sprintf("localhost:%d", *callbackPort), + Handler: mux, + } + if err := r.server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) { + r.errChan <- err + } +} + +// getAuthorizationCode implements the AuthorizationCodeFetcher interface. +// It displays the authorization URL to the user and waits for the callback. +func (r *codeReceiver) getAuthorizationCode(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + fmt.Printf("\nPlease open the following URL in your browser to authenticate:\n%s\n\n", args.URL) + select { + case authRes := <-r.authChan: + return authRes, nil + case err := <-r.errChan: + return nil, err + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +// close shuts down the HTTP server. +func (r *codeReceiver) close() { + if r.server != nil { + r.server.Close() + } +} + +func main() { + flag.Parse() + + // Validate required configuration. + if *idpClientID == "" { + log.Fatal("--idp_client_id is required") + } + + // Set up the OAuth callback receiver. + receiver := &codeReceiver{ + authChan: make(chan *auth.AuthorizationResult), + errChan: make(chan error), + } + listener, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", *callbackPort)) + if err != nil { + log.Fatalf("failed to listen on port %d: %v", *callbackPort, err) + } + go receiver.serveRedirectHandler(listener) + defer receiver.close() + + log.Printf("OAuth callback server listening on http://localhost:%d", *callbackPort) + + // Create an ID Token fetcher that performs OIDC login with the enterprise IdP. + idTokenFetcher := func(ctx context.Context) (*oauth2.Token, error) { + log.Println("Starting OIDC login flow...") + + creds := &oauthex.ClientCredentials{ + ClientID: *idpClientID, + } + if *idpClientSecret != "" { + creds.ClientSecretAuth = &oauthex.ClientSecretAuth{ + ClientSecret: *idpClientSecret, + } + } + + oidcConfig := &extauth.OIDCLoginConfig{ + IssuerURL: *idpIssuerURL, + Credentials: creds, + RedirectURL: fmt.Sprintf("http://localhost:%d", *callbackPort), + Scopes: []string{"openid", "profile", "email"}, + } + + // PerformOIDCLogin handles the complete OIDC Authorization Code flow with PKCE. + tokens, err := extauth.PerformOIDCLogin(ctx, oidcConfig, receiver.getAuthorizationCode) + if err != nil { + return nil, fmt.Errorf("OIDC login failed: %w", err) + } + + log.Println("OIDC login successful, obtained ID token") + return tokens, nil + } + + // Create the Enterprise Handler. + // This handler implements the complete Enterprise Managed Authorization flow: + // 1. OIDC Login: User authenticates with enterprise IdP → ID Token (via idTokenFetcher). + // 2. Token Exchange (RFC 8693): ID Token → ID-JAG at IdP. + // 3. JWT Bearer Grant (RFC 7523): ID-JAG → Access Token at MCP Server. + log.Println("Creating enterprise authorization handler...") + + // Prepare IdP credentials + idpCreds := &oauthex.ClientCredentials{ + ClientID: *idpClientID, + } + if *idpClientSecret != "" { + idpCreds.ClientSecretAuth = &oauthex.ClientSecretAuth{ + ClientSecret: *idpClientSecret, + } + } + + // Prepare MCP credentials + var mcpCreds *oauthex.ClientCredentials + if *mcpClientID != "" { + mcpCreds = &oauthex.ClientCredentials{ + ClientID: *mcpClientID, + } + if *mcpClientSecret != "" { + mcpCreds.ClientSecretAuth = &oauthex.ClientSecretAuth{ + ClientSecret: *mcpClientSecret, + } + } + } + + enterpriseHandler, err := extauth.NewEnterpriseHandler(&extauth.EnterpriseHandlerConfig{ + // IdP configuration (where the user authenticates). + IdPIssuerURL: *idpIssuerURL, + IdPCredentials: idpCreds, + + // MCP Server configuration (the resource being accessed). + MCPAuthServerURL: *mcpAuthServerURL, + MCPResourceURI: *mcpResourceURI, + MCPCredentials: mcpCreds, + MCPScopes: []string{"read", "write"}, + + // ID Token fetcher (performs OIDC login when needed). + IDTokenFetcher: idTokenFetcher, + }) + if err != nil { + log.Fatalf("failed to create enterprise handler: %v", err) + } + + // Create the MCP client transport with the enterprise handler. + transport := &mcp.StreamableClientTransport{ + Endpoint: *mcpServerURL, + OAuthHandler: enterpriseHandler, + } + + ctx := context.Background() + client := mcp.NewClient(&mcp.Implementation{ + Name: "enterprise-client-example", + Version: "1.0.0", + }, nil) + + log.Printf("Connecting to MCP server at %s...", *mcpServerURL) + session, err := client.Connect(ctx, transport, nil) + if err != nil { + log.Fatalf("failed to connect to MCP server: %v", err) + } + defer session.Close() + + log.Println("Successfully connected to MCP server!") + + // List available tools as a demonstration. + tools, err := session.ListTools(ctx, nil) + if err != nil { + log.Fatalf("failed to list tools: %v", err) + } + + log.Println("\nAvailable tools:") + if len(tools.Tools) == 0 { + log.Println(" (no tools available)") + } else { + for _, tool := range tools.Tools { + log.Printf(" - %q: %s", tool.Name, tool.Description) + } + } +} diff --git a/examples/auth/server/main.go b/examples/auth/server/main.go new file mode 100644 index 0000000..94ad9ae --- /dev/null +++ b/examples/auth/server/main.go @@ -0,0 +1,167 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package main + +import ( + "context" + "encoding/json" + "flag" + "fmt" + "log" + "net/http" + "net/http/httputil" + "net/url" + "strings" + "time" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/oauthex" +) + +// Flags. +var ( + port = flag.Int("port", 8000, "Port to listen on") +) + +// Configuration required for this example. +var ( + // Authorization server to return in the protected resource metadata. + authorizationServer = "" + // Introspection endpoint for verifying tokens. + introspectionEndpoint = "" + // Client credentials used in the introspection request. + clientID = "" + clientSecret = "" +) + +func verifyToken(ctx context.Context, token string, _ *http.Request) (*auth.TokenInfo, error) { + data := url.Values{} + data.Set("token", token) + data.Set("token_type_hint", "access_token") + + req, err := http.NewRequestWithContext(ctx, "POST", introspectionEndpoint, strings.NewReader(data.Encode())) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("Accept", "application/json") + req.SetBasicAuth(clientID, clientSecret) + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + dump, _ := httputil.DumpResponse(resp, true) + log.Printf("Introspection failed: %s", dump) + return nil, fmt.Errorf("introspection failed with status %d", resp.StatusCode) + } + + var result struct { + Active bool `json:"active"` + Scope string `json:"scope"` + Exp int64 `json:"exp"` + Sub string `json:"sub"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, err + } + + if !result.Active { + return nil, auth.ErrInvalidToken + } + + return &auth.TokenInfo{ + Scopes: strings.Fields(result.Scope), + Expiration: time.Unix(result.Exp, 0), + UserID: result.Sub, + }, nil +} + +type args struct { + Input string `json:"input"` +} + +func echo(ctx context.Context, req *mcp.CallToolRequest, args args) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: args.Input}, + }, + }, nil, nil +} + +func main() { + flag.Parse() + metadata := &oauthex.ProtectedResourceMetadata{ + Resource: fmt.Sprintf("http://localhost:%d/mcp", *port), + AuthorizationServers: []string{authorizationServer}, + ScopesSupported: []string{"read"}, + } + http.Handle("/.well-known/oauth-protected-resource", auth.ProtectedResourceMetadataHandler(metadata)) + + server := mcp.NewServer(&mcp.Implementation{ + Name: "test-server", + Version: "1.0.0", + }, nil) + server.AddReceivingMiddleware(createLoggingMiddleware()) + mcp.AddTool(server, &mcp.Tool{Name: "echo"}, echo) + + handler := mcp.NewStreamableHTTPHandler(func(req *http.Request) *mcp.Server { + return server + }, nil) + + authMiddleware := auth.RequireBearerToken(verifyToken, &auth.RequireBearerTokenOptions{ + Scopes: []string{"read"}, + ResourceMetadataURL: fmt.Sprintf("http://localhost:%d/.well-known/oauth-protected-resource", *port), + }) + + http.Handle("/mcp", authMiddleware(handler)) + + log.Printf("Starting server on http://localhost:%d", *port) + log.Fatal(http.ListenAndServe(fmt.Sprintf("localhost:%d", *port), nil)) +} + +// createLoggingMiddleware creates an MCP middleware that logs method calls. +func createLoggingMiddleware() mcp.Middleware { + return func(next mcp.MethodHandler) mcp.MethodHandler { + return func( + ctx context.Context, + method string, + req mcp.Request, + ) (mcp.Result, error) { + start := time.Now() + sessionID := req.GetSession().ID() + + // Log request details. + log.Printf("[REQUEST] Session: %s | Method: %s", + sessionID, + method) + + // Call the actual handler. + result, err := next(ctx, method, req) + + // Log response details. + duration := time.Since(start) + + if err != nil { + log.Printf("[RESPONSE] Session: %s | Method: %s | Status: ERROR | Duration: %v | Error: %v", + sessionID, + method, + duration, + err) + } else { + log.Printf("[RESPONSE] Session: %s | Method: %s | Status: OK | Duration: %v", + sessionID, + method, + duration) + } + + return result, err + } + } +} diff --git a/examples/client/listfeatures/main.go b/examples/client/listfeatures/main.go index 044b5e9..193d29f 100644 --- a/examples/client/listfeatures/main.go +++ b/examples/client/listfeatures/main.go @@ -39,7 +39,7 @@ func main() { fmt.Fprintln(os.Stderr, "Usage: listfeatures --http=\"https://example.com/server/mcp\"") fmt.Fprintln(os.Stderr, "List all features for a stdio MCP server") fmt.Fprintln(os.Stderr) - fmt.Fprintln(os.Stderr, "Example:\n\tlistfeatures npx @modelcontextprotocol/server-everything") + fmt.Fprintln(os.Stderr, "Example:\n\tlistfeatures go run github.com/modelcontextprotocol/go-sdk/examples/server/hello") os.Exit(2) } diff --git a/examples/client/loadtest/main.go b/examples/client/loadtest/main.go index cbf7cd6..31b74a6 100644 --- a/examples/client/loadtest/main.go +++ b/examples/client/loadtest/main.go @@ -69,9 +69,7 @@ func main() { // Run the test. var wg sync.WaitGroup for range *workers { - wg.Add(1) - go func() { - defer wg.Done() + wg.Go(func() { client := mcp.NewClient(&mcp.Implementation{Name: "mcp-client", Version: "v1.0.0"}, nil) cs, err := client.Connect(parentCtx, &mcp.StreamableClientTransport{Endpoint: args[0]}, nil) if err != nil { @@ -108,7 +106,7 @@ func main() { } } } - }() + }) } wg.Wait() stop() // restore the interrupt signal diff --git a/examples/http/README.md b/examples/http/README.md index 3bd8c93..ae6bccb 100644 --- a/examples/http/README.md +++ b/examples/http/README.md @@ -15,7 +15,7 @@ The example implements: ```bash go run . server ``` -This starts an MCP server on `http://localhost:8080` (default) that provides a `cityTime` tool. +This starts an MCP server on `http://localhost:8000` (default) that provides a `cityTime` tool. To run a client in another terminal: diff --git a/examples/http/logging_middleware.go b/examples/http/logging_middleware.go index 4266012..97d8ad3 100644 --- a/examples/http/logging_middleware.go +++ b/examples/http/logging_middleware.go @@ -5,47 +5,49 @@ package main import ( + "context" "log" - "net/http" "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" ) -// responseWriter wraps http.ResponseWriter to capture the status code. -type responseWriter struct { - http.ResponseWriter - statusCode int -} +// createLoggingMiddleware creates an MCP middleware that logs method calls. +func createLoggingMiddleware() mcp.Middleware { + return func(next mcp.MethodHandler) mcp.MethodHandler { + return func( + ctx context.Context, + method string, + req mcp.Request, + ) (mcp.Result, error) { + start := time.Now() + sessionID := req.GetSession().ID() -func (rw *responseWriter) WriteHeader(code int) { - rw.statusCode = code - rw.ResponseWriter.WriteHeader(code) -} + // Log request details. + log.Printf("[REQUEST] Session: %s | Method: %s", + sessionID, + method) -func loggingHandler(handler http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - start := time.Now() + // Call the actual handler. + result, err := next(ctx, method, req) - // Create a response writer wrapper to capture status code. - wrapped := &responseWriter{ResponseWriter: w, statusCode: http.StatusOK} + // Log response details. + duration := time.Since(start) - // Log request details. - log.Printf("[REQUEST] %s | %s | %s %s", - start.Format(time.RFC3339), - r.RemoteAddr, - r.Method, - r.URL.Path) + if err != nil { + log.Printf("[RESPONSE] Session: %s | Method: %s | Status: ERROR | Duration: %v | Error: %v", + sessionID, + method, + duration, + err) + } else { + log.Printf("[RESPONSE] Session: %s | Method: %s | Status: OK | Duration: %v", + sessionID, + method, + duration) + } - // Call the actual handler. - handler.ServeHTTP(wrapped, r) - - // Log response details. - duration := time.Since(start) - log.Printf("[RESPONSE] %s | %s | %s %s | Status: %d | Duration: %v", - time.Now().Format(time.RFC3339), - r.RemoteAddr, - r.Method, - r.URL.Path, - wrapped.statusCode, - duration) - }) + return result, err + } + } } diff --git a/examples/http/main.go b/examples/http/main.go index 188674a..ced4871 100644 --- a/examples/http/main.go +++ b/examples/http/main.go @@ -117,6 +117,9 @@ func runServer(url string) { Version: "1.0.0", }, nil) + // Add MCP-level logging middleware. + server.AddReceivingMiddleware(createLoggingMiddleware()) + // Add the cityTime tool. mcp.AddTool(server, &mcp.Tool{ Name: "cityTime", @@ -128,13 +131,11 @@ func runServer(url string) { return server }, nil) - handlerWithLogging := loggingHandler(handler) - log.Printf("MCP server listening on %s", url) log.Printf("Available tool: cityTime (cities: nyc, sf, boston)") - // Start the HTTP server with logging handler. - if err := http.ListenAndServe(url, handlerWithLogging); err != nil { + // Start the HTTP server. + if err := http.ListenAndServe(url, handler); err != nil { log.Fatalf("Server failed: %v", err) } } diff --git a/examples/server/auth-middleware/README.md b/examples/server/auth-middleware/README.md index aaae767..001237c 100644 --- a/examples/server/auth-middleware/README.md +++ b/examples/server/auth-middleware/README.md @@ -44,6 +44,7 @@ go test -cover ### Public Endpoints (No Authentication Required) - `GET /health` - Health check +- `GET /.well-known/oauth-protected-resource` - OAuth 2.0 Protected Resource Metadata (RFC 9728) ### MCP Endpoints (Authentication Required) @@ -228,6 +229,41 @@ authenticatedHandler := authMiddleware(customMiddleware(mcpHandler)) 4. **Token Expiration**: Set appropriate token expiration times 5. **Principle of Least Privilege**: Grant only the minimum required scopes +## CORS Support + +The OAuth metadata endpoint (`/.well-known/oauth-protected-resource`) supports CORS to enable +cross-origin client discovery. It sets `Access-Control-Allow-Origin: *` by default because +OAuth metadata is public information meant for client discovery (RFC 9728 §3.1). + +### Custom CORS Policies + +For more sophisticated CORS requirements (origin validation, credentials, etc.), wrap the handler +with a CORS middleware library: + +**Using github.com/rs/cors:** +```go +import "github.com/rs/cors" + +c := cors.New(cors.Options{ + AllowedOrigins: []string{"https://example.com"}, + AllowedMethods: []string{"GET", "OPTIONS"}, +}) +http.Handle("/.well-known/oauth-protected-resource", + c.Handler(auth.ProtectedResourceMetadataHandler(metadata))) +``` + +**Using github.com/jub0bs/cors:** +```go +import "github.com/jub0bs/cors" + +corsMiddleware, err := cors.NewMiddleware(cors.Config{ + Origins: []string{"https://example.com"}, + Methods: []string{"GET", "OPTIONS"}, +}) +http.Handle("/.well-known/oauth-protected-resource", + corsMiddleware.Wrap(auth.ProtectedResourceMetadataHandler(metadata))) +``` + ## Use Cases **Ideal for:** diff --git a/examples/server/auth-middleware/go.mod b/examples/server/auth-middleware/go.mod index cbc89c7..f1bfc37 100644 --- a/examples/server/auth-middleware/go.mod +++ b/examples/server/auth-middleware/go.mod @@ -1,15 +1,16 @@ module auth-middleware-example -go 1.23.0 +go 1.25.0 require ( - github.com/golang-jwt/jwt/v5 v5.2.2 + github.com/golang-jwt/jwt/v5 v5.3.1 github.com/modelcontextprotocol/go-sdk v0.3.0 ) require ( - github.com/google/jsonschema-go v0.3.0 // indirect + github.com/google/jsonschema-go v0.4.2 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + golang.org/x/oauth2 v0.35.0 // indirect ) replace github.com/modelcontextprotocol/go-sdk => ../../../ diff --git a/examples/server/auth-middleware/go.sum b/examples/server/auth-middleware/go.sum index 7b7a8e5..a14b134 100644 --- a/examples/server/auth-middleware/go.sum +++ b/examples/server/auth-middleware/go.sum @@ -1,10 +1,15 @@ github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8= github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= +github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q= github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI= +golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= +golang.org/x/oauth2 v0.35.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA= golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo= golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg= diff --git a/examples/server/auth-middleware/main.go b/examples/server/auth-middleware/main.go index 75a38b8..f9ac34f 100644 --- a/examples/server/auth-middleware/main.go +++ b/examples/server/auth-middleware/main.go @@ -20,12 +20,15 @@ import ( "github.com/golang-jwt/jwt/v5" "github.com/modelcontextprotocol/go-sdk/auth" "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/oauthex" ) // This example demonstrates how to integrate auth.RequireBearerToken middleware // with an MCP server to provide authenticated access to MCP tools and resources. -var httpAddr = flag.String("http", ":8080", "HTTP address to listen on") +// Default to loopback to avoid exposing powerful demo endpoints to the local network. +// Set -http 0.0.0.0:8080 to intentionally expose it. +var httpAddr = flag.String("http", "localhost:8080", "HTTP address to listen on") // JWTClaims represents the claims in our JWT tokens. // In a real application, you would include additional claims like issuer, audience, etc. @@ -237,7 +240,8 @@ func main() { // Create authentication middleware. jwtAuth := auth.RequireBearerToken(verifyJWT, &auth.RequireBearerTokenOptions{ - Scopes: []string{"read"}, // Require "read" permission + Scopes: []string{"read"}, // Require "read" permission + ResourceMetadataURL: "http://localhost:8080/.well-known/oauth-protected-resource", }) apiKeyAuth := auth.RequireBearerToken(verifyAPIKey, &auth.RequireBearerTokenOptions{ @@ -340,22 +344,36 @@ func main() { }) }) + // OAuth protected resource metadata endpoint. + // This endpoint provides OAuth configuration information to clients. + // CORS is enabled by default to support cross-origin client discovery. + metadata := &oauthex.ProtectedResourceMetadata{ + Resource: "http://localhost:8080/mcp/jwt", + AuthorizationServers: []string{ + "https://auth.example.com/.well-known/openid-configuration", + }, + ScopesSupported: []string{"read", "write"}, + } + http.Handle("/.well-known/oauth-protected-resource", + auth.ProtectedResourceMetadataHandler(metadata)) + // Start the HTTP server. log.Println("Authenticated MCP Server") log.Println("========================") log.Println("Server starting on", *httpAddr) log.Println() log.Println("Available endpoints:") - log.Println(" GET /health - Health check (no auth)") - log.Println(" GET /generate-token - Generate JWT token") - log.Println(" POST /generate-api-key - Generate API key") - log.Println(" POST /mcp/jwt - MCP endpoint (JWT auth)") - log.Println(" POST /mcp/apikey - MCP endpoint (API key auth)") + log.Println(" GET /health - Health check (no auth)") + log.Println(" GET /.well-known/oauth-protected-resource - OAuth resource metadata") + log.Println(" GET /generate-token - Generate JWT token") + log.Println(" POST /generate-api-key - Generate API key") + log.Println(" POST /mcp/jwt - MCP endpoint (JWT auth)") + log.Println(" POST /mcp/apikey - MCP endpoint (API key auth)") log.Println() log.Println("Available MCP Tools:") - log.Println(" - say_hi - Simple greeting (any auth)") - log.Println(" - get_user_info - Get user info (read scope)") - log.Println(" - create_resource - Create resource (write scope)") + log.Println(" - say_hi - Simple greeting (any auth)") + log.Println(" - get_user_info - Get user info (read scope)") + log.Println(" - create_resource - Create resource (write scope)") log.Println() log.Println("Example usage:") log.Println(" # Generate a token") diff --git a/examples/server/custom-method/main.go b/examples/server/custom-method/main.go new file mode 100644 index 0000000..6b2861f --- /dev/null +++ b/examples/server/custom-method/main.go @@ -0,0 +1,110 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +// The custom-method example demonstrates registering and calling a custom +// JSON-RPC method that is not part of the standard MCP spec. +// +// The server registers a "latin/translate" method that translates simple +// English phrases into Latin. A client connects over an in-memory transport, +// calls the custom method, and prints the result. +package main + +import ( + "context" + "fmt" + "log" + "strings" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +type TranslateParams struct { + mcp.ParamsBase + Text string `json:"text"` +} + +type TranslateResult struct { + mcp.ResultBase + Latin string `json:"latin"` +} + +var translations = map[string]string{ + "hello": "salve", + "goodbye": "vale", + "thank you": "gratias tibi ago", + "how are you": "quid agis", + "good morning": "bonum mane", + "good night": "bonam noctem", + "friend": "amicus", + "water": "aqua", + "love": "amor", + "war": "bellum", + "peace": "pax", + "truth": "veritas", + "light": "lux", + "time": "tempus", + "life": "vita", + "death": "mors", + "star": "stella", + "earth": "terra", + "sea": "mare", + "the die is cast": "alea iacta est", + "i came i saw i conquered": "veni vidi vici", + "seize the day": "carpe diem", +} + +func main() { + ctx := context.Background() + + server := mcp.NewServer(&mcp.Implementation{Name: "latin-server", Version: "v1.0.0"}, nil) + + if err := mcp.AddReceivingCustomMethod(server, "latin/translate", + func(ctx context.Context, ss *mcp.ServerSession, params *TranslateParams) (*TranslateResult, error) { + key := strings.ToLower(strings.TrimSpace(params.Text)) + latin, ok := translations[key] + if !ok { + latin = fmt.Sprintf("[unknown: %q — try: %s]", params.Text, knownPhrases()) + } + return &TranslateResult{Latin: latin}, nil + }); err != nil { + log.Fatal(err) + } + + ct, st := mcp.NewInMemoryTransports() + + ss, err := server.Connect(ctx, st, nil) + if err != nil { + log.Fatal(err) + } + defer ss.Close() + + client := mcp.NewClient(&mcp.Implementation{Name: "latin-client", Version: "v1.0.0"}, nil) + if err := mcp.AddSendingCustomMethod[*TranslateParams, *TranslateResult](client, "latin/translate"); err != nil { + log.Fatal(err) + } + + cs, err := client.Connect(ctx, ct, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + phrases := []string{"Hello", "Seize the day", "Peace", "Truth", "I came I saw I conquered"} + for _, phrase := range phrases { + result, err := mcp.CallCustomMethod[*TranslateParams, *TranslateResult]( + ctx, cs, "latin/translate", &TranslateParams{Text: phrase}) + if err != nil { + log.Fatalf("translate %q: %v", phrase, err) + } + fmt.Printf("%-35s → %s\n", phrase, result.Latin) + } +} + +func knownPhrases() string { + phrases := make([]string, 0, len(translations)) + for k := range translations { + phrases = append(phrases, fmt.Sprintf("%q", k)) + } + return strings.Join(phrases, ", ") +} diff --git a/examples/server/distributed/main.go b/examples/server/distributed/main.go index b144040..888cba5 100644 --- a/examples/server/distributed/main.go +++ b/examples/server/distributed/main.go @@ -16,6 +16,8 @@ // Example: // // ./distributed -http=localhost:8080 -child_ports=8081,8082 +// +//lint:file-ignore SA1019 example server exercises deprecated SEP-2577 logging API for demonstration. package main import ( @@ -89,16 +91,14 @@ func parent() { cmd.Env = append(os.Environ(), fmt.Sprintf("%s=%s", childPortVar, port)) cmd.Stderr = os.Stderr - wg.Add(1) - go func() { - defer wg.Done() + wg.Go(func() { log.Printf("starting child %d at %s", i, childURL) if err := cmd.Run(); err != nil && ctx.Err() == nil { log.Printf("child %d failed: %v", i, err) } else { log.Printf("child %d exited gracefully", i) } - }() + }) } // Start a reverse proxy that round-robin's requests to each backend. @@ -116,13 +116,11 @@ func parent() { }, } - wg.Add(1) - go func() { - defer wg.Done() + wg.Go(func() { if err := server.ListenAndServe(); err != nil && ctx.Err() == nil { log.Printf("Server failed: %v", err) } - }() + }) log.Printf("Serving at %s (CTRL-C to cancel)", *httpAddr) diff --git a/examples/server/everything/main.go b/examples/server/everything/main.go index 0b81919..a08cac8 100644 --- a/examples/server/everything/main.go +++ b/examples/server/everything/main.go @@ -3,10 +3,14 @@ // license that can be found in the LICENSE file. // The everything server implements all supported features of an MCP server. +// +//lint:file-ignore SA1019 example server exercises deprecated SEP-2577 APIs (roots, sampling, logging) for demonstration. package main import ( "context" + _ "embed" + "encoding/base64" "flag" "fmt" "log" @@ -16,6 +20,7 @@ import ( "os" "runtime" "strings" + "sync/atomic" "github.com/google/jsonschema-go/jsonschema" "github.com/modelcontextprotocol/go-sdk/mcp" @@ -50,19 +55,25 @@ func main() { CompletionHandler: complete, // support completions by setting this handler } - server := mcp.NewServer(&mcp.Implementation{Name: "everything"}, opts) + // Add an icon to the server implementation. + icons := mcpIcons() + server := mcp.NewServer(&mcp.Implementation{Name: "everything", WebsiteURL: "https://example.com", Icons: icons}, opts) // Add tools that exercise different features of the protocol. mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, contentTool) - mcp.AddTool(server, &mcp.Tool{Name: "greet (structured)"}, structuredTool) // returns structured output - mcp.AddTool(server, &mcp.Tool{Name: "ping"}, pingingTool) // performs a ping - mcp.AddTool(server, &mcp.Tool{Name: "log"}, loggingTool) // performs a log - mcp.AddTool(server, &mcp.Tool{Name: "sample"}, samplingTool) // performs sampling - mcp.AddTool(server, &mcp.Tool{Name: "elicit"}, elicitingTool) // performs elicitation - mcp.AddTool(server, &mcp.Tool{Name: "roots"}, rootsTool) // lists roots + mcp.AddTool(server, &mcp.Tool{Name: "greet (structured)"}, structuredTool) // returns structured output + mcp.AddTool(server, &mcp.Tool{Name: "greet (with Icons)", Icons: icons}, structuredTool) // tool with icons + mcp.AddTool(server, &mcp.Tool{Name: "greet (content with ResourceLink)"}, resourceLinkContentTool(icons)) // tool that returns content with a resource link + mcp.AddTool(server, &mcp.Tool{Name: "ping"}, pingingTool) // performs a ping + mcp.AddTool(server, &mcp.Tool{Name: "log"}, loggingTool) // performs a log + mcp.AddTool(server, &mcp.Tool{Name: "sample"}, samplingTool) // performs sampling + mcp.AddTool(server, &mcp.Tool{Name: "elicit (form)"}, elicitFormTool) // performs form elicitation + mcp.AddTool(server, &mcp.Tool{Name: "elicit (url)"}, elicitURLTool) // performs url elicitation + mcp.AddTool(server, &mcp.Tool{Name: "roots"}, rootsTool) // lists roots // Add a basic prompt. server.AddPrompt(&mcp.Prompt{Name: "greet"}, prompt) + server.AddPrompt(&mcp.Prompt{Name: "greet (with Icons)", Icons: icons}, prompt) // greet prompt with icons // Add an embedded resource. server.AddResource(&mcp.Resource{ @@ -70,6 +81,20 @@ func main() { MIMEType: "text/plain", URI: "embedded:info", }, embeddedResource) + server.AddResource(&mcp.Resource{ // text resource with icons + Name: "info (with Icons)", + MIMEType: "text/plain", + URI: "embedded:info", + Icons: icons, + }, embeddedResource) + + // Add a resource template. + server.AddResourceTemplate(&mcp.ResourceTemplate{ + Name: "Resource template (with Icon)", + MIMEType: "text/plain", + URITemplate: "http://example.com/~{resource_name}/", + Icons: icons, + }, embeddedResource) // Serve over stdio, or streamable HTTP if -http is set. if *httpAddr != "" { @@ -140,6 +165,23 @@ func contentTool(ctx context.Context, req *mcp.CallToolRequest, args args) (*mcp }, nil, nil } +// resourceLinkContentTool returns a ResourceLink content with icons. +func resourceLinkContentTool(icons []mcp.Icon) func(ctx context.Context, req *mcp.CallToolRequest, args args) (*mcp.CallToolResult, any, error) { + return func(ctx context.Context, req *mcp.CallToolRequest, args args) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.ResourceLink{ + Name: "greeting", + Title: "A friendly greeting", + MIMEType: "text/plain", + URI: "data:text/plain,Hi%20" + url.PathEscape(args.Name), + Icons: icons, + }, + }, + }, nil, nil + } +} + type result struct { Message string `json:"message" jsonschema:"the message to convey"` } @@ -194,7 +236,7 @@ func samplingTool(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.Ca }, nil, nil } -func elicitingTool(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { +func elicitFormTool(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { res, err := req.Session.Elicit(ctx, &mcp.ElicitParams{ Message: "provide a random string", RequestedSchema: &jsonschema.Schema{ @@ -214,6 +256,26 @@ func elicitingTool(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.C }, nil, nil } +var elicitations atomic.Int32 + +func elicitURLTool(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { + elicitID := fmt.Sprintf("%d", elicitations.Add(1)) + _, err := req.Session.Elicit(ctx, &mcp.ElicitParams{ + Message: "submit a string", + URL: fmt.Sprintf("http://localhost:6062?id=%s", elicitID), + ElicitationID: elicitID, + }) + if err != nil { + return nil, nil, fmt.Errorf("eliciting failed: %v", err) + } + // TODO: actually wait for the elicitation form to be submitted. + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "(elicitation pending)"}, + }, + }, nil, nil +} + func complete(ctx context.Context, req *mcp.CompleteRequest) (*mcp.CompleteResult, error) { return &mcp.CompleteResult{ Completion: mcp.CompletionResultDetails{ @@ -222,3 +284,15 @@ func complete(ctx context.Context, req *mcp.CompleteRequest) (*mcp.CompleteResul }, }, nil } + +//go:embed mcp.png +var mcpIconData []byte + +func mcpIcons() []mcp.Icon { + return []mcp.Icon{{ + Source: "data:image/png;base64," + base64.StdEncoding.EncodeToString(mcpIconData), + MIMEType: "image/png", + Sizes: []string{"48x48"}, + Theme: mcp.IconThemeLight, + }} +} diff --git a/examples/server/everything/mcp.png b/examples/server/everything/mcp.png new file mode 100644 index 0000000..8e08571 Binary files /dev/null and b/examples/server/everything/mcp.png differ diff --git a/examples/server/memory/main.go b/examples/server/memory/main.go index 99d109a..039c21b 100644 --- a/examples/server/memory/main.go +++ b/examples/server/memory/main.go @@ -61,7 +61,7 @@ type DeleteEntitiesArgs struct { // DeleteObservationsArgs defines the delete observations tool parameters. type DeleteObservationsArgs struct { - Deletions []Observation `json:"deletions" mcp:"obeservations to delete"` + Deletions []Observation `json:"deletions" mcp:"observations to delete"` } // DeleteRelationsArgs defines the delete relations tool parameters. diff --git a/examples/server/middleware/main.go b/examples/server/middleware/main.go index 4e37471..262d726 100644 --- a/examples/server/middleware/main.go +++ b/examples/server/middleware/main.go @@ -38,7 +38,7 @@ func main() { logger.Info("MCP method started", "method", method, "session_id", req.GetSession().ID(), - "has_params", req.GetParams() != nil, + "has_params", mcp.HasParams(req), ) // Log more for tool calls. if ctr, ok := req.(*mcp.CallToolRequest); ok { diff --git a/examples/server/proxy/README.md b/examples/server/proxy/README.md new file mode 100644 index 0000000..5aa37c1 --- /dev/null +++ b/examples/server/proxy/README.md @@ -0,0 +1,52 @@ +# Proxy Header Propagation Example + +This example demonstrates how to propagate HTTP headers from an incoming MCP server request to an outgoing MCP client request. + +This is particularly useful when building a gateway or proxy server that needs to forward authentication tokens, trace IDs, or other context-sensitive headers to downstream services. + +## Architecture + +This example runs three components in a single process: + +1. **Backend Server** (Port 8082): + - Exposes a tool `echo_headers` that returns the headers it received. + +2. **Proxy Server** (Port 8081): + - Exposes a tool `forward_headers`. + - Acts as a client to the Backend Server. + - Uses a custom `http.RoundTripper` (`HeaderForwardingTransport`) to inject headers from the context into outgoing requests. + +3. **Client**: + - Connects to the Proxy Server. + - Calls `forward_headers`. + +## How it works + +1. The Client calls `forward_headers` on the Proxy Server. +2. The Proxy Server receives the request. The request context contains the HTTP headers in `req.Extra.Header`. +3. The Proxy's tool handler extracts these headers and places them into a new `context.Context` using a specific key (`headerContextKey`). +4. The Proxy uses an `mcp.Client` configured with a custom `HTTPClient` that uses `HeaderForwardingTransport`. +5. `HeaderForwardingTransport.RoundTrip` inspects the context of outgoing requests. If it finds headers under `headerContextKey`, it adds them to the HTTP request. +6. The Backend Server receives the request with the propagated headers. + +## Running the Example + +Run the example with: + +```bash +go run main.go +``` + +You should see output indicating: +1. Backend and Proxy servers starting. +2. Gateway receiving headers. +3. Client receiving the result, which contains the echoed headers. + +Example output: + +``` +2025/08/29 10:00:00 Starting Backend Server on :8082 +2025/08/29 10:00:00 Starting Gateway Server on :8081 +2025/08/29 10:00:00 Gateway received headers: map[...] +2025/08/29 10:00:00 Client received result: ... +``` diff --git a/examples/server/proxy/main.go b/examples/server/proxy/main.go new file mode 100644 index 0000000..3a4469b --- /dev/null +++ b/examples/server/proxy/main.go @@ -0,0 +1,141 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package main + +import ( + "context" + "fmt" + "log" + "net/http" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// Define a context key for passing headers. +type contextKey string + +const headerContextKey contextKey = "ctx-headers" + +// HeaderForwardingTransport is an http.RoundTripper that injects headers +// from the context into the outgoing request. +type HeaderForwardingTransport struct{} + +// RoundTrip executes a single HTTP transaction, adding headers from the context. +func (h *HeaderForwardingTransport) RoundTrip(req *http.Request) (*http.Response, error) { + // Clone the request to avoid modifying the original one. + newReq := req.Clone(req.Context()) + + // Retrieve headers from the context. + if headers, ok := req.Context().Value(headerContextKey).(http.Header); ok { + for key, values := range headers { + // Copy all values for each header. + for _, value := range values { + newReq.Header.Add(key, value) + } + } + } + + return http.DefaultTransport.RoundTrip(newReq) +} + +// NewHeaderForwardingClient creates a new http.Client that uses HeaderForwardingTransport. +func NewHeaderForwardingClient() *http.Client { + return &http.Client{ + Transport: new(HeaderForwardingTransport), + } +} + +// Backend Server: Echoes received headers to verify propagation. +func runBackendServer() { + server := mcp.NewServer(&mcp.Implementation{Name: "backend-server", Version: "1.0.0"}, nil) + + mcp.AddTool(server, &mcp.Tool{ + Name: "echo_headers", + Description: "Returns the headers received by the server", + }, func(ctx context.Context, req *mcp.CallToolRequest, args struct{}) (*mcp.CallToolResult, any, error) { + return nil, req.Extra.Header, nil + }) + + // Start the backend server on port 8082. + log.Println("Starting Backend Server on :8082") + if err := http.ListenAndServe(":8082", mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { + return server + }, nil)); err != nil { + log.Fatal(err) + } +} + +// Proxy Server: Forwards requests to the backend with headers. +func runProxyServer(ctx context.Context) { + // Connect to the backend server (acting as a client). + client := mcp.NewClient(&mcp.Implementation{Name: "proxy-client", Version: "1.0.0"}, nil) + clientSession, err := client.Connect(ctx, &mcp.StreamableClientTransport{ + Endpoint: "http://localhost:8082/mcp", + HTTPClient: NewHeaderForwardingClient(), + }, nil) + if err != nil { + log.Fatalf("Failed to connect to backend: %v", err) + } + defer clientSession.Close() + + server := mcp.NewServer(&mcp.Implementation{Name: "proxy-server", Version: "1.0.0"}, nil) + // Add a tool that forwards the call to the backend. + mcp.AddTool(server, &mcp.Tool{ + Name: "forward_headers", + Description: "Calls the backend server, ensuring headers are propagated", + }, func(ctx context.Context, req *mcp.CallToolRequest, args struct{}) (*mcp.CallToolResult, any, error) { + incomingHeaders := req.Extra.Header + log.Printf("Gateway received headers: %v", incomingHeaders) + propagateCtx := context.WithValue(ctx, headerContextKey, incomingHeaders) + result, err := clientSession.CallTool(propagateCtx, &mcp.CallToolParams{ + Name: "echo_headers", + }) + if err != nil { + return nil, nil, fmt.Errorf("failed to call backend: %w", err) + } + + return result, nil, nil + }) + + // Start the gateway server on port 8081. + log.Println("Starting Gateway Server on :8081") + if err := http.ListenAndServe(":8081", mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { + return server + }, nil)); err != nil { + log.Fatal(err) + } +} + +func runClient(ctx context.Context) { + // Connect to the proxy server (acting as a client). + client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "1.0.0"}, nil) + clientSession, err := client.Connect(ctx, &mcp.StreamableClientTransport{ + Endpoint: "http://localhost:8081/mcp", + }, nil) + if err != nil { + log.Fatalf("Failed to connect to proxy: %v", err) + } + defer clientSession.Close() + + result, err := clientSession.CallTool(ctx, &mcp.CallToolParams{ + Name: "forward_headers", + }) + if err != nil { + log.Fatalf("Failed to call proxy: %v", err) + } + log.Printf("Client received result: %v", result.StructuredContent) +} + +func main() { + ctx := context.Background() + go runBackendServer() + // Give the backend a moment to start. + time.Sleep(100 * time.Millisecond) + go runProxyServer(ctx) + // Give the proxy a moment to start. + time.Sleep(100 * time.Millisecond) + runClient(ctx) +} diff --git a/examples/server/rate-limiting/go.mod b/examples/server/rate-limiting/go.mod index adf535b..dc37554 100644 --- a/examples/server/rate-limiting/go.mod +++ b/examples/server/rate-limiting/go.mod @@ -1,6 +1,6 @@ module github.com/modelcontextprotocol/go-sdk/examples/rate-limiting -go 1.23.0 +go 1.25.0 require ( github.com/modelcontextprotocol/go-sdk v0.3.0 @@ -8,7 +8,7 @@ require ( ) require ( - github.com/google/jsonschema-go v0.3.0 // indirect + github.com/google/jsonschema-go v0.4.2 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect ) diff --git a/examples/server/rate-limiting/go.sum b/examples/server/rate-limiting/go.sum index 92c2739..cb67bcb 100644 --- a/examples/server/rate-limiting/go.sum +++ b/examples/server/rate-limiting/go.sum @@ -2,6 +2,7 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q= github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE= diff --git a/examples/server/sequentialthinking/main.go b/examples/server/sequentialthinking/main.go index e0ae521..4775c33 100644 --- a/examples/server/sequentialthinking/main.go +++ b/examples/server/sequentialthinking/main.go @@ -234,7 +234,7 @@ func deepCopyThoughts(thoughts []*Thought) []*Thought { func StartThinking(ctx context.Context, _ *mcp.CallToolRequest, args StartThinkingArgs) (*mcp.CallToolResult, any, error) { sessionID := args.SessionID if sessionID == "" { - sessionID = randText() + sessionID = rand.Text() } estimatedSteps := args.EstimatedSteps @@ -480,20 +480,6 @@ func ThinkingHistory(ctx context.Context, req *mcp.ReadResourceRequest) (*mcp.Re }, nil } -// Copied from crypto/rand. -// TODO: once 1.24 is assured, just use crypto/rand. -const base32alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567" - -func randText() string { - // ⌈log₃₂ 2¹²⁸⌉ = 26 chars - src := make([]byte, 26) - rand.Read(src) - for i := range src { - src[i] = base32alphabet[src[i]%32] - } - return string(src) -} - func main() { flag.Parse() diff --git a/examples/server/toolschemas/main.go b/examples/server/toolschemas/main.go index abc1e3f..0a2cfbd 100644 --- a/examples/server/toolschemas/main.go +++ b/examples/server/toolschemas/main.go @@ -36,7 +36,7 @@ func simpleGreeting(_ context.Context, _ *mcp.CallToolRequest, input Input) (*mc // manualGreeter handles the parsing and validation of input and output manually. // -// Therfore, it needs to close over its resolved schemas, to use them in +// Therefore, it needs to close over its resolved schemas, to use them in // validation. type manualGreeter struct { inputSchema *jsonschema.Resolved diff --git a/go.mod b/go.mod index d5917a1..3287a95 100644 --- a/go.mod +++ b/go.mod @@ -1,12 +1,20 @@ module github.com/modelcontextprotocol/go-sdk -go 1.23.0 +go 1.25.0 require ( - github.com/golang-jwt/jwt/v5 v5.2.2 + github.com/golang-jwt/jwt/v5 v5.3.1 github.com/google/go-cmp v0.7.0 - github.com/google/jsonschema-go v0.3.0 + github.com/google/jsonschema-go v0.4.3 + github.com/segmentio/encoding v0.5.4 github.com/yosida95/uritemplate/v3 v3.0.2 - golang.org/x/oauth2 v0.30.0 - golang.org/x/tools v0.34.0 + golang.org/x/oauth2 v0.35.0 + golang.org/x/time v0.15.0 + golang.org/x/tools v0.42.0 +) + +require ( + github.com/segmentio/asm v1.1.3 // indirect + golang.org/x/sync v0.20.0 // indirect + golang.org/x/sys v0.41.0 // indirect ) diff --git a/go.sum b/go.sum index 32ceedf..c13454a 100644 --- a/go.sum +++ b/go.sum @@ -1,12 +1,22 @@ -github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8= -github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= +github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= +github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= -github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q= -github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/google/jsonschema-go v0.4.3 h1:/DBOLZTfDow7pe2GmaJNhltueGTtDKICi8V8p+DQPd0= +github.com/google/jsonschema-go v0.4.3/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc= +github.com/segmentio/asm v1.1.3/go.mod h1:Ld3L4ZXGNcSLRg4JBsZ3//1+f/TjYl0Mzen/DQy1EJg= +github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0= +github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= -golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI= -golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= -golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo= -golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg= +golang.org/x/oauth2 v0.35.0 h1:Mv2mzuHuZuY2+bkyWXIHMfhNdJAdwW3FuWeCPYN5GVQ= +golang.org/x/oauth2 v0.35.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k= +golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= +golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k= +golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0= diff --git a/internal/authutil/scopes.go b/internal/authutil/scopes.go new file mode 100644 index 0000000..4b5de73 --- /dev/null +++ b/internal/authutil/scopes.go @@ -0,0 +1,36 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. +package authutil + +import ( + "maps" + "slices" + "strings" + + "golang.org/x/oauth2" +) + +// UnionScopes returns the union of the existing and challenged scope sets. +// It is used during step-up authorization to accumulate scopes across +// authorization rounds (SEP-2350). +func UnionScopes(existing, challenged []string) []string { + combined := make(map[string]struct{}) + for _, s := range existing { + combined[s] = struct{}{} + } + for _, s := range challenged { + combined[s] = struct{}{} + } + return slices.Collect(maps.Keys(combined)) +} + +// ScopesFromToken extracts the granted scopes from an OAuth2 token response. +// Per RFC 6749 §5.1, the scope parameter is optional; returns nil if absent. +func ScopesFromToken(token *oauth2.Token) []string { + scope, ok := token.Extra("scope").(string) + if !ok { + return nil + } + return strings.Fields(scope) +} diff --git a/internal/authutil/scopes_test.go b/internal/authutil/scopes_test.go new file mode 100644 index 0000000..1fdfd7d --- /dev/null +++ b/internal/authutil/scopes_test.go @@ -0,0 +1,72 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. +package authutil + +import ( + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" +) + +func TestUnionScopes(t *testing.T) { + tests := []struct { + name string + existing []string + challenged []string + want []string + }{ + { + name: "both empty", + existing: nil, + challenged: nil, + want: nil, + }, + { + name: "existing only", + existing: []string{"read"}, + challenged: nil, + want: []string{"read"}, + }, + { + name: "challenged only", + existing: nil, + challenged: []string{"write"}, + want: []string{"write"}, + }, + { + name: "disjoint scopes", + existing: []string{"read"}, + challenged: []string{"write"}, + want: []string{"read", "write"}, + }, + { + name: "overlapping scopes", + existing: []string{"read", "write"}, + challenged: []string{"write", "admin"}, + want: []string{"read", "write", "admin"}, + }, + { + name: "identical scopes", + existing: []string{"read", "write"}, + challenged: []string{"read", "write"}, + want: []string{"read", "write"}, + }, + { + name: "mixed scopes", + existing: []string{"b", "a"}, + challenged: []string{"c", "a"}, + want: []string{"a", "b", "c"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := UnionScopes(tt.existing, tt.challenged) + if diff := cmp.Diff(tt.want, got, cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" { + t.Errorf("UnionScopes() mismatch (-want +got):\n%s", diff) + } + }) + } +} diff --git a/internal/authutil/util.go b/internal/authutil/util.go new file mode 100644 index 0000000..880d1bf --- /dev/null +++ b/internal/authutil/util.go @@ -0,0 +1,13 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package authutil + +import "strings" + +// IssuersEqual reports whether two OAuth 2.0 authorization server issuer +// identifiers refer to the same server comparing them without the final trailing slash. +func IssuersEqual(a, b string) bool { + return strings.TrimSuffix(a, "/") == strings.TrimSuffix(b, "/") +} diff --git a/internal/authutil/util_test.go b/internal/authutil/util_test.go new file mode 100644 index 0000000..9d3ff7a --- /dev/null +++ b/internal/authutil/util_test.go @@ -0,0 +1,29 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package authutil + +import "testing" + +func TestIssuersEqual(t *testing.T) { + tests := []struct { + a, b string + want bool + }{ + {"https://issuer.example.com", "https://issuer.example.com", true}, + {"https://issuer.example.com/", "https://issuer.example.com", true}, + {"https://issuer.example.com", "https://issuer.example.com/", true}, + {"https://issuer.example.com/", "https://issuer.example.com/", true}, + {"https://issuer.example.com/tenant", "https://issuer.example.com/tenant", true}, + {"https://issuer.example.com/tenant/", "https://issuer.example.com/tenant", true}, + {"https://issuer.example.com", "https://other.example.com", false}, + {"https://issuer.example.com/a", "https://issuer.example.com/b", false}, + {"", "", true}, + } + for _, tt := range tests { + if got := IssuersEqual(tt.a, tt.b); got != tt.want { + t.Errorf("IssuersEqual(%q, %q) = %v, want %v", tt.a, tt.b, got, tt.want) + } + } +} diff --git a/internal/docs/README.src.md b/internal/docs/README.src.md index 7f7b224..3283efa 100644 --- a/internal/docs/README.src.md +++ b/internal/docs/README.src.md @@ -1,4 +1,21 @@ -# Features +# Overview + +The SDK consists of several importable packages: + +- The + [`github.com/modelcontextprotocol/go-sdk/mcp`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp) + package defines the primary APIs for constructing and using MCP clients and + servers. +- The + [`github.com/modelcontextprotocol/go-sdk/jsonrpc`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/jsonrpc) package is for users implementing + their own transports. +- The + [`github.com/modelcontextprotocol/go-sdk/auth`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth) + package provides some primitives for supporting OAuth. +- The + [`github.com/modelcontextprotocol/go-sdk/oauthex`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/oauthex) + package provides extensions to the OAuth protocol, such as ProtectedResourceMetadata. + These docs mirror the [official MCP spec](https://modelcontextprotocol.io/specification/2025-06-18). Use the index below to learn how the SDK implements a particular aspect of the @@ -37,3 +54,14 @@ protocol. # TroubleShooting See [troubleshooting.md](troubleshooting.md) for a troubleshooting guide. + +# Backwards compatibility + +See [mcpgodebug.md](mcpgodebug.md) for a list of backwards incompatible behavior changes +and description how they can be temporarily undone. + +# Rough edges + +See [rough_edges.md](rough_edges.md) for a list of rough edges or API +oversights that can't be addressed due to our compatibility promise. We'll +revisit these if/when we move to a v2 of the SDK. diff --git a/internal/docs/client.src.md b/internal/docs/client.src.md index fc37d45..c0f64b4 100644 --- a/internal/docs/client.src.md +++ b/internal/docs/client.src.md @@ -4,6 +4,13 @@ ## Roots +> **Note:** The roots feature is deprecated as of protocol version 2026-07-28 +> ([SEP-2577](https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging)). +> It remains fully functional during the deprecation window (at least twelve +> months). The SDK continues to support roots for compatibility. New code +> should pass paths via tool parameters, resource URIs, or configuration +> instead. + MCP allows clients to specify a set of filesystem ["roots"](https://modelcontextprotocol.io/specification/2025-06-18/client/roots). The SDK supports this as follows: @@ -21,11 +28,32 @@ client, a call to `AddRoot` or `RemoveRoots` will result in a [`ServerSession.ListRoots`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession.ListRoots) method. To receive notifications about root changes, set [`ServerOptions.RootsListChangedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.RootsListChangedHandler). +For protocol versions `2026-07-28` and later, `ListRoots` requests are +delivered via the +[Multi Round-Trip Requests](#multi-round-trip-requests) +pattern. %include ../../mcp/client_example_test.go roots - +### Roots list changed + +`Client.AddRoots` and `Client.RemoveRoots` notify every connected server that +the list changed. Servers observe this through +[`ServerOptions.RootsListChangedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.RootsListChangedHandler); +as with the server-side list-changed notifications, it reports only that +something changed, so read the list back with `ServerSession.ListRoots`. + +%include ../../mcp/client_example_test.go rootslistchanged - + ## Sampling +> **Note:** The sampling feature is deprecated as of protocol version +> 2026-07-28 +> ([SEP-2577](https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging)). +> It remains fully functional during the deprecation window (at least twelve +> months). The SDK continues to support sampling for compatibility. Servers +> that need LLM completions should call LLM provider APIs directly. + [Sampling](https://modelcontextprotocol.io/specification/2025-06-18/client/sampling) is a way for servers to leverage the client's AI capabilities. It is implemented in the SDK as follows: @@ -37,6 +65,11 @@ This function is invoked whenever the server requests sampling. **Server-side**: To use sampling from the server, call [`ServerSession.CreateMessage`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession.CreateMessage). +For protocol versions `2026-07-28` and later, sampling requests are +delivered via the +[Multi Round-Trip Requests](#multi-round-trip-requests) +pattern. + %include ../../mcp/client_example_test.go sampling - ## Elicitation @@ -47,9 +80,158 @@ allows servers to request user inputs. It is implemented in the SDK as follows: **Client-side**: To add the `elicitation` capability to a client, set [`ClientOptions.ElicitationHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.ElicitationHandler). The elicitation handler must return a result that matches the requested schema; -otherwise, elicitation returns an error. +otherwise, elicitation returns an error. If your handler supports [URL mode +elicitation](https://modelcontextprotocol.io/specification/2025-11-25/client/elicitation#url-mode-elicitation-requests), +you must declare that capability explicitly (see [Capabilities](#capabilities)) **Server-side**: To use elicitation from the server, call [`ServerSession.Elicit`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession.Elicit). +For protocol versions `2026-07-28` and later, elicitation requests are +delivered via the +[Multi Round-Trip Requests](#multi-round-trip-requests) +pattern. + %include ../../mcp/client_example_test.go elicitation - + +### Schema defaults and enums + +`ElicitParams.RequestedSchema` is a flat schema of primitive fields, which the +client renders as a form. Two field keywords shape that form. + +A `Default` ([SEP-1034](https://modelcontextprotocol.io/seps/1034--support-default-values-for-all-primitive-types-in)) prefills a +field. When the user accepts without supplying it, the SDK fills the field in +from the schema before the result reaches either side's caller — the client +does so after its elicitation handler returns, and `ServerSession.Elicit` does +so again on receipt. This is unconditional; there is no opt-in flag. Marking a +defaulted field `Required` defeats it: accepted content is validated against +the schema before defaults are applied, so an answer that omits the field is +rejected rather than defaulted. + +An `Enum` ([SEP-1330](https://modelcontextprotocol.io/seps/1330-elicitation-enum-schema-improvements-and-standards)) restricts a +field to a fixed set of values, which the client renders as a choice. Enums +are supported only on `"string"` fields; declaring one on another type is +rejected. To label the choices, set the legacy `enumNames` keyword through +`Schema.Extra`, with exactly one name per enum value — a mismatched length is +rejected. + +%include ../../mcp/client_example_test.go elicitationschema - + +### Completing a URL elicitation + +In URL mode the user finishes out of band, in a browser, so nothing in the +elicitation result tells the client when they are done. The server signals that +with +[`ServerSession.NotifyElicitationComplete`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession.NotifyElicitationComplete), +passing the same `ElicitationID` the request carried; the client observes it +through +[`ClientOptions.ElicitationCompleteHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.ElicitationCompleteHandler). +Send it from whatever endpoint the hosted flow redirects back to. + +The notification matters most when a handler rejects a request with +[`URLElicitationRequiredError`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#URLElicitationRequiredError): +the client parks the original request until a notification names that +`ElicitationID`, and then retries it automatically. Until one arrives, the +client waits. + +%include ../../mcp/client_example_test.go elicitationcomplete - + +## Multi Round-Trip Requests + +[SEP-2322](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2322) +introduces the MRTR pattern: server-to-client requests for sampling, +elicitation, and roots are no longer issued as fresh JSON-RPC requests but +are carried inside the in-flight reply of a `tools/call`, `prompts/get`, or +`resources/read`. The client must respond by retrying the original request +with the produced responses. + +The SDK installs `clientMultiRoundTripMiddleware` for every client by +default. The middleware: + +1. Inspects each `tools/call`/`prompts/get`/`resources/read` reply. +2. If the result's `NeedsInput()` is true, fans out the `InputRequests` map + concurrently, calling the configured handler for each (`elicit`, + `createMessage`/`createMessageWithTools`, or `listRoots`). +3. Threads the server-supplied opaque `RequestState` back unchanged. +4. Retries the original request with the responses set, repeating until the + result no longer needs input. + +The middleware is enabled by default. To opt out, set +[`ClientOptions.MultiRoundTrip.Disabled = true`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#MultiRoundTripOptions); +the client will then surface input-required results directly to the caller +(the returned `CallToolResult`, `GetPromptResult`, or `ReadResourceResult` +will report `NeedsInput() == true` and expose the server's `InputRequests` +and opaque `RequestState`). Your code must fulfil each request and re-issue +the original call with `InputResponses` set and `RequestState` echoed +back. + +For legacy (`<= 2025-11-25`) servers, the SDK transparently sends server +requests on the legacy server-initiated channel; the MRTR machinery is a +no-op in that direction. For legacy clients talking to MRTR-style servers, +the server SDK applies the inverse compatibility shim — see the +[server-side documentation](server.md#multi-round-trip-requests). + +## Capabilities + +Client capabilities are advertised to servers during the initialization +handshake. [By default](rough_edges.md), the SDK advertises the `logging` +capability. Additional capabilities are automatically added when server +features are added (e.g. via `AddTool`), or when handlers are set in the +`ServerOptions` struct (e.g., setting `CompletionHandler` adds the +`completions` capability), or may be configured explicitly. + +### Capability inference + +When handlers are set on `ClientOptions` (e.g., `CreateMessageHandler` or +`ElicitationHandler`), the corresponding capability is automatically added if +not already present, with a default configuration. + +For elicitation, if the handler is set but no `Capabilities.Elicitation` is +specified, the client defaults to form elicitation. To enable URL elicitation +or both modes, [configure `Capabilities.Elicitation` +explicitly](#explicit-capabilities). + +See the [`ClientCapabilities` +documentation](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientCapabilities) +for further details on inference. + +### Explicit capabilities + +To explicitly declare capabilities, or to override the [default inferred +capability](#capability-inference), set +[`ClientOptions.Capabilities`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.Capabilities). +This sets the initial client capabilities, before any capabilities are added +based on configured handlers. If a capability is already present in +`Capabilities`, adding a handler will not change its configuration. + +This allows you to: + +- **Disable default capabilities**: Pass an empty `&ClientCapabilities{}` to + disable all defaults, including roots. +- **Disable listChanged notifications**: Set `ListChanged: false` on a + capability to prevent the client from sending list-changed notifications + when roots are added or removed. +- **Configure elicitation modes**: Specify which elicitation modes (form, URL) + the client supports. + +```go +// Configure elicitation modes and disable roots. +client := mcp.NewClient(impl, &mcp.ClientOptions{ + Capabilities: &mcp.ClientCapabilities{ + Elicitation: &mcp.ElicitationCapabilities{ + Form: &mcp.FormElicitationCapabilities{}, + URL: &mcp.URLElicitationCapabilities{}, + }, + }, + ElicitationHandler: handler, +}) +``` + +### Extensions + +[SEP-2133](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2133) +adds an `extensions` map to `ClientCapabilities` and `ServerCapabilities` so +that optional capabilities outside the core protocol can be declared on the +wire. Keys are namespaced as `"{vendor-prefix}/{extension-name}"`; values +are per-extension settings objects. + diff --git a/internal/docs/doc.go b/internal/docs/doc.go index 7b23ad6..4164060 100644 --- a/internal/docs/doc.go +++ b/internal/docs/doc.go @@ -8,6 +8,9 @@ //go:generate weave -o ../../docs/client.md ./client.src.md //go:generate weave -o ../../docs/server.md ./server.src.md //go:generate weave -o ../../docs/troubleshooting.md ./troubleshooting.src.md +//go:generate weave -o ../../docs/rough_edges.md ./rough_edges.src.md +//go:generate weave -o ../../docs/mcpgodebug.md ./mcpgodebug.src.md +//go:generate weave -o ../../docs/quick_start.md ./quick_start.src.md // The doc package generates the documentation at /doc, via go:generate. // diff --git a/internal/docs/mcpgodebug.src.md b/internal/docs/mcpgodebug.src.md new file mode 100644 index 0000000..8d1385f --- /dev/null +++ b/internal/docs/mcpgodebug.src.md @@ -0,0 +1,177 @@ +# Backwards compatibility and MCPGODEBUG + + According to our compatibility promise, we can't break backward compatibility + of the SDK API. However, sometimes we need to change the behavior of the SDK + in a backward-incompatible way in order to fix bugs or security issues. + In those cases we introduce temporary compatibility parameters, that can be + used to opt-out of the new behavior. They are usually maintained for two + minor release cycles and then removed. + + The compatibility parameters are provided via the `MCPGODEBUG` environment + variable. The value of the variable is a comma-separated list of parameter + value assignments, e.g.: + + ``` + MCPGODEBUG=parameter1=value1,parameter2=value2 + ``` + +## `MCPGODEBUG` history + +### 1.9.0 + +Options listed below were added and will be removed in the 1.11.0 version of the SDK. + +- `structuredcontentfloat64` added. If set to `1`, numbers in + `CallToolResult.StructuredContent` are decoded as `float64`, restoring the + previous behavior. The default behavior decodes numbers as `json.Number` so + that large integers and other exact JSON values do not silently lose + precision. + +### 1.8.0 + +Options listed below were added and will be removed in the 1.9.0 version of the SDK. + +- `plaintextstatefulrejection` added. If set to `1`, a stateful + `StreamableHTTPHandler` will respond with a plain-text `http.Error` 400 body + when it receives a request carrying per-request metadata (i.e. an + `io.modelcontextprotocol/protocolVersion` `_meta` field, or an + `MCP-Protocol-Version` header >= `2026-07-28`), restoring the previous + behavior. The default behavior was changed so that the server responds with a + JSON-RPC error of code `CodeUnsupportedProtocolVersion` (`-32022`) carrying + an `UnsupportedProtocolVersionData` payload that advertises the legacy + versions the server supports. This lets the client's + existing renegotiation logic recover and prevents the failure from tearing + down the underlying connection. + +- `blockingcancelnotify` added. If set to `1`, a cancelled call waits + synchronously for the best-effort `notifications/cancelled` message to be + delivered (up to `notifyCancellationTimeout`, currently 5 s) before + returning to the caller, restoring the previous behavior. The delivery + error is joined into the caller's returned error. The default behavior was + changed so that the call is retired immediately and the notification is + sent asynchronously off the caller's return path: the caller returns as + soon as its context is cancelled and cannot be delayed by a slow or + unresponsive peer. See issue #1150. + +Options below were removed, according to plan: + +- `seterroroverwrite`. `SetError` now always preserves existing `Content`. + +- `enableoriginverification`. A nil + `StreamableHTTPOptions.CrossOriginProtection` now always means no + cross-origin protection. Wrap the handler with cross-origin protection + middleware instead. + +- `disablecontenttypecheck`. Content-Type validation on HTTP POST requests can + no longer be disabled. + +- `disablelocalhostprotection`, whose removal had been postponed to this + release. Set the `DisableLocalhostProtection` field in + `StreamableHTTPOptions` or `SSEOptions` to opt out of DNS rebinding + protection. + +### 1.7.0 + +Options listed below were added and will be removed in the 1.9.0 version of the SDK. + +- `customresnotfounderrcode` added. If set to `1`, `ResourceNotFoundError` will + use the custom error code `-32002` instead of the standard `-32602` (Invalid + Params), restoring the previous behavior. The default behavior was changed to + align with SEP-2164 and the JSON-RPC specification. + +- `hintomitempty` added. If set to `1`, `ToolAnnotations` JSON marshaling + will omit `ReadOnlyHint` and `IdempotentHint` when their value is `false`, + restoring the previous behavior. The default behavior was changed to always + serialize these fields, since their Go types are bare `bool` (not `*bool`) + and omitting `false` made it indistinguishable from unset. + +- `allowsessionsinstateless` added. If set to `1`, stateless streamable HTTP + servers will read the `Mcp-Session-Id` request header (or generate one via + `GetSessionID`), set it on response headers, and accept `DELETE` requests, + restoring the previous behavior. The default behavior was changed so that + stateless servers ignore session IDs entirely and reject `DELETE` with 405. + +- `nomethodnotfoundcodeinerror` added. If set to `1`, the jsonrpc2 layer will not + include the MethodNotFound Error (`-32601`) in the error response when the + requested method in STDIO transport is not found. The default behavior was + changed to include the MethodNotFound Error in the error response when the + requested method in STDIO transport is not found. + +- `noprotocolerrorbody` added. If set to `1`, the streamable HTTP client will + not attempt to decode the JSON-RPC error body of a non-2xx HTTP response, + and any non-transient error will permanently fail the connection, restoring + the previous behavior. The default behavior was changed so that the client + attempts to decode the JSON-RPC error body of a non-2xx response, surfaces + the underlying JSON-RPC error, and wraps it with `jsonrpc2.ErrRejected` so + that per-call rejections do not tear down the session. + +- `nowrapinvalidparams` added. If set to `1`, the server will not wrap + params-decoding failures with `jsonrpc2.ErrInvalidParams`, so wire responses + carry the zero-value error code `0` instead of `-32602` ("invalid params"), + restoring the previous behavior. The default behavior was changed so that + the server always reports params-decoding failures with the standard + `-32602` code. + +- `disablecompleteparamsvalidation` added. If set to `1`, `Server.complete` + will not validate that the required `ref` and `argument.name` fields on + `CompleteParams` are present, restoring the previous behavior of dispatching + the request to the completion handler unconditionally. The default behavior + was changed to reject malformed requests with `-32602` (Invalid Params). + +### 1.6.1 + +Options listed below were added and will be removed in the 1.8.0 version of the SDK. + +- `disablecontenttypecheck` added. If set to `1`, Content-Type validation on + HTTP POST requests will be disabled, allowing requests with non-JSON or missing + Content-Type headers. The default behavior is to validate that HTTP POST + requests have Content-Type: application/json. + +### 1.6.0 + +Options listed below were added and will be removed in the 1.8.0 version of the SDK. + +- `seterroroverwrite` added. If set to `1`, `SetError` will always overwrite + `Content` with the error text, restoring the previous behavior. The default + behavior was changed to preserve existing `Content` if it has already been + populated. + +- `enableoriginverification` added. If set to `1`, default (zero-value) + cross-origin protection will be applied when + `StreamableHTTPOptions.CrossOriginProtection` is nil, restoring the + behavior from v1.4.1-v1.5.0. The default behavior was changed to not + enable cross-origin protection. + +- `disablelocalhostprotection` removal was postponed until 1.8.0, as it is now + also used in the SSE transport. + +Options below were removed: + +- `jsonescaping`, according to plan, + +- `disablecrossoriginprotection`, it was replaced by + `enableoriginverification` after the default was changed to not enable + cross-origin protection. + +### 1.4.1 + +Options listed below will be removed in the 1.6.0 version of the SDK. + +- `disablecrossoriginprotection` added. If set to `1`, newly added cross-origin + protection will be disabled. The default behavior was changed to enable + cross-origin protection. + +### 1.4.0 + +Options listed below will be removed in the 1.6.0 version of the SDK. + +- `jsonescaping` added. If set to `1`, JSON marshaling will preserve the previous + behavior of escaping HTML characters in JSON strings. The default behavior + was changed to not escape HTML characters, to be consistent with other SDKs. + +- `disablelocalhostprotection` added. If set to `1`, newly added DNS rebinding + protection will be disabled. The default behavior was changed to enable DNS rebinding + protection. The protection can also be disabled by setting the + `DisableLocalhostProtection` field in the `StreamableHTTPOptions` or + `SSEOptions` struct to `true`, which is the recommended way to disable + the protection long term. **Removal of this option was postponed until 1.8.0.** diff --git a/internal/docs/protocol.src.md b/internal/docs/protocol.src.md index e85374a..7d6b01b 100644 --- a/internal/docs/protocol.src.md +++ b/internal/docs/protocol.src.md @@ -6,15 +6,23 @@ The SDK provides an API for defining both MCP clients and servers, and connecting them over various transports. When a client and server are -connected, it creates a logical session, which follows the MCP spec's -[lifecycle](https://modelcontextprotocol.io/specification/2025-06-18/basic/lifecycle). +connected, it creates a logical session. -In this SDK, both a -[`Client`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Client) -and -[`Server`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server) -can handle multiple peers. Every time a new peer is connected, it creates a new -session. +The MCP specification defines two lifecycle models, and the SDK supports both +transparently based on the negotiated protocol version: + +- The **legacy `initialize` handshake** + ([lifecycle](https://modelcontextprotocol.io/specification/2025-06-18/basic/lifecycle)) + used by protocol versions through `2025-11-25`. +- A **stateless** model introduced in `2026-07-28` by + [SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575), + in which there is no `initialize`/`notifications/initialized` handshake, and + each request carries its protocol version and client capabilities in + `_meta`. Clients SHOULD also include their identity (`clientInfo`) on every + request, and servers SHOULD include their identity (`serverInfo`) on every + response. + +In both models, the SDK exposes the same API: - A `Client` is a logical MCP client, configured with various [`ClientOptions`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions). @@ -30,11 +38,16 @@ session. [`Server.Connect`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server.Connect), it creates a [`ServerSession`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession). - This session is not initialized until the client sends the - `notifications/initialized` message. Use `ServerOptions.InitializedHandler` - to listen for this event, or just use the session through various feature - handlers (such as a `ToolHandler`). Requests to the server are rejected until - the client has initialized the session. + +In the legacy model, the server session is not considered initialized until +the client sends the `notifications/initialized` message. Use +`ServerOptions.InitializedHandler` to listen for this event, or just use the +session through various feature handlers (such as a `ToolHandler`). Requests +to the server are rejected until the client has initialized the session. + +In the stateless model (`2026-07-28`+), there is no handshake. The server +processes the first request the moment it arrives, validates the per-request +`_meta` fields. Both `ClientSession` and `ServerSession` have a `Close` method to terminate the session, and a `Wait` method to await session termination by the peer. Typically, @@ -42,6 +55,82 @@ it is the client's responsibility to end the session. %include ../../mcp/mcp_example_test.go lifecycle - +### Discovery + +Introduced in `2026-07-28` by +[SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575), +the `server/discover` RPC lets a client discover a server's supported +protocol versions, capabilities, and identity before issuing any other +request. Servers implementing `2026-07-28` MUST implement it. + +- **Server**: `Server.Connect` registers the + `server/discover` handler automatically; the response is computed from the + server's static configuration (capabilities, server info, instructions) and + the transport-filtered list of supported protocol versions. +- **Client**: `Client.Connect` calls `server/discover` first and + uses the result to negotiate a mutually supported version. If discovery + fails or the server does not support the latest version, the client falls back to the + legacy `initialize` handshake. + +A server advertises and negotiates every protocol version the SDK supports; +set [`ServerOptions.SupportedProtocolVersions`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions) +to narrow that set, for example to stop serving a revision your deployment has +retired. + +```go +server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v1.0.0"}, &mcp.ServerOptions{ + SupportedProtocolVersions: []string{"2026-07-28", "2025-11-25"}, +}) +``` + +The set can only narrow support, never widen it; the field's documentation +covers how each protocol era answers a request at an excluded version. + +### Per-request metadata keys + +When the negotiated protocol version is `2026-07-28` or later, every request +carries these keys inside its `_meta` map (constants live in +[`mcp/protocol.go`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#pkg-constants)): + +| Constant | Wire key | Type | Required | +|---|---|---|---| +| `MetaKeyProtocolVersion` | `io.modelcontextprotocol/protocolVersion` | `string` | Yes | +| `MetaKeyClientCapabilities` | `io.modelcontextprotocol/clientCapabilities` | `*ClientCapabilities` | Yes | +| `MetaKeyClientInfo` | `io.modelcontextprotocol/clientInfo` | `*Implementation` | No | +| `MetaKeyLogLevel` | `io.modelcontextprotocol/logLevel` | `LoggingLevel` (deprecated by SEP-2577) | No | + +The client populates the required keys automatically on every outgoing +request, and populates `clientInfo` when configured with an `*Implementation` +(the default). Server-side handlers can read them through +`ServerRequest[P].ProtocolVersion()`, `ServerRequest[P].ClientInfo()`, and +`ServerRequest[P].ClientCapabilities()`. + +### Per-response metadata keys + +Under the same protocol version, servers SHOULD identify themselves on every +response. The SDK populates this +key automatically on every outgoing response: + +| Constant | Wire key | Type | Required | +|---|---|---|---| +| `MetaKeyServerInfo` | `io.modelcontextprotocol/serverInfo` | `*Implementation` | No | + +### Subscriptions + +Introduced in `2026-07-28` by +[SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575), +`subscriptions/listen` replaces the legacy `resources/subscribe` RPC and +the GET-based SSE endpoint with a single long-lived request that +multiplexes every kind of server-to-client change notification a client +opts in to. On the wire the client sends one `subscriptions/listen` +request whose `notifications` field enumerates what it wants +(`toolsListChanged`, `promptsListChanged`, `resourcesListChanged`, and/or +a list of resource URIs in `resourceSubscriptions`); the server replies +first with a `notifications/subscriptions/acknowledged` notification +reporting the honored subset, then streams every change notification on +the same request, and finally closes with a `SubscriptionsListenResult` +when it tears the subscription down. + ## Transports A @@ -106,12 +195,66 @@ connect to the server: transport := &mcp.StreamableClientTransport{ Endpoint: "http://localhost:8080/mcp", } -client, err := mcp.Connect(ctx, transport, &mcp.ClientOptions{...}) +client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v1.0.0"}, nil) +session, err := client.Connect(ctx, transport, nil) ``` The `StreamableClientTransport` handles the HTTP requests and communicates with the server using the streamable transport protocol. +#### HTTP Headers + +[SEP-2243](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2243) +standardises a small set of MCP-specific HTTP headers, sent and validated by +the SDK only for sessions negotiated on protocol version `2026-07-28` or +later. + +| Header | Direction | Purpose | +|---|---|---| +| `Mcp-Protocol-Version` | request | Mirrors `_meta.io.modelcontextprotocol/protocolVersion` from the body | +| `Mcp-Session-Id` | request/response | Logical session identifier (removed in stateless mode) | +| `Mcp-Method` | request | Mirrors the JSON-RPC `method` field; mismatch ⇒ `-32020` | +| `Mcp-Name` | request | Mirrors the request's principal name (`tools/call.params.name`, `prompts/get.params.name`, `resources/read.params.uri`); mismatch ⇒ `-32020` | +| `Mcp-Param-{Header}` | request | Per-tool parameter passthrough; see below | + +The header-name set is exhaustive: for protocol versions earlier than +`2026-07-28`, only `Mcp-Protocol-Version` and `Mcp-Session-Id` are recognised. + +**Body↔header mirroring.** When `Mcp-Method` or `Mcp-Name` is present but +does not match the JSON-RPC body, the server returns +[`CodeHeaderMismatch`](#error-codes) (`-32020`). The same applies to a +mismatch between `Mcp-Protocol-Version` and the body's +`_meta.io.modelcontextprotocol/protocolVersion`. + +**`Mcp-Param-*` passthrough.** A tool may annotate properties of its +`InputSchema` with `x-mcp-header: ""`. When a client invokes that +tool over HTTP, the SDK serializes the matching argument(s) as +`Mcp-Param-` request headers (using `=?base64?...?=` encoding for +non-ASCII values). On the server, the headers are validated against the body +and either accepted or rejected with `-32020`. Properties so annotated must +be typed `string`, `integer`, or `boolean`; tools that fail validation at +registration time are silently dropped from `tools/list`. + +#### Resumability and Redelivery + +By default, the streamable server does not support [resumability or +redelivery](https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#resumability-and-redelivery) +of messages, because doing so requires either a persistent storage solution or +unbounded memory usage (see also +[#580](https://github.com/modelcontextprotocol/go-sdk/issues/580)). + +To enable resumability, set `StreamableHTTPOptions.EventStore` to a non-nil +value. The SDK provides a `MemoryEventStore` for testing or simple use cases; +for production use it is generally advisable to use a more sophisticated +implementation. + +> **Note**: [SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575) +> removes SSE stream resumability (`Last-Event-ID`, SSE event IDs) for protocol +> version `2026-07-28`. The SDK preserves the `EventStore` and `Last-Event-ID` +> code paths for backward compatibility with `2025-11-25` and earlier; on +> `2026-07-28` sessions, a broken response stream loses the in-flight request +> and the client must re-issue it as a new request with a new ID. + #### Stateless Mode The streamable server supports a _stateless mode_ by setting @@ -130,10 +273,42 @@ to see the logical session > modelcontextprotocol/modelcontextprotocol#1372, or > modelcontextprotocol/modelcontextprotocol#1442 for potential refinements. -_See [examples/server/distributed](../examples/server/distributed/main.go) for -an example using statless mode to implement a server distributed across +> **Required for `2026-07-28`**: the streamable HTTP transport accepts +> requests at protocol version `2026-07-28` **only** when `Stateless = true`. +> Requests at that version against a non-stateless handler are rejected. + +_See [examples/server/distributed](https://github.com/modelcontextprotocol/go-sdk/blob/main/examples/server/distributed/main.go) for +an example using stateless mode to implement a server distributed across multiple processes._ +### Legacy SSE Transport + +Before the streamable transport, the +[2024-11-05](https://modelcontextprotocol.io/specification/2024-11-05/basic/transports) +spec defined an HTTP transport built from two endpoints: a hanging GET that +streams server-to-client messages as server-sent events, and a per-session +endpoint the client POSTs its messages to. The SDK still implements it, for +talking to peers that predate the streamable transport. New deployments should +use the [streamable transport](#streamable-transport) instead. + +It is implemented across three types, mirroring the streamable ones: + +- [`SSEHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#SSEHandler): + an `http.Handler` that creates a session per incoming GET and routes POSTs to + the right session. Construct it with + [`NewSSEHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#NewSSEHandler), + which takes a function returning the `Server` to serve a given request. +- [`SSEServerTransport`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#SSEServerTransport): + one such session. Reach for it directly only when serving the two endpoints + yourself; `SSEHandler` creates these for you. +- [`SSEClientTransport`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#SSEClientTransport): + the client side, which needs only the `Endpoint` of the GET. + +%include ../../mcp/sse_example_test.go ssehandler - + +_See [examples/server/sse](https://github.com/modelcontextprotocol/go-sdk/blob/main/examples/server/sse/main.go) +for a standalone server._ + ### Custom transports The SDK supports [custom @@ -142,7 +317,7 @@ by implementing the [`Transport`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Transport) interface: a logical bidirectional stream of JSON-RPC messages. -_Full example: [examples/server/custom-transport](../examples/server/custom-transport/main.go)._ +_Full example: [examples/server/custom-transport](https://github.com/modelcontextprotocol/go-sdk/blob/main//examples/server/custom-transport/main.go)._ ### Concurrency @@ -176,7 +351,7 @@ and invokes the passed to `RequireBearerToken` to parse the token and perform validation. The middleware function checks expiration and scopes (if they are provided in [`RequireBearerTokenOptions.Scopes`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#RequireBearerTokenOptions.Scopes)), so the -`TokenVerifer` doesn't have to. +`TokenVerifier` doesn't have to. If [`RequireBearerTokenOptions.ResourceMetadataURL`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#RequireBearerTokenOptions.ResourceMetadataURL) is set and verification fails, the middleware function sets the WWW-Authenticate header as required by the [Protected Resource Metadata spec](https://datatracker.ietf.org/doc/html/rfc9728). @@ -186,38 +361,201 @@ from `req.Extra.TokenInfo`, where `req` is the handler's request. (For example, [`CallToolRequest`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#CallToolRequest).) HTTP handlers wrapped by the `RequireBearerToken` middleware can obtain the `TokenInfo` from the context with [`auth.TokenInfoFromContext`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#TokenInfoFromContext). - + +#### OAuth Protected Resource Metadata + +Servers implementing OAuth 2.0 authorization should expose a protected resource metadata endpoint +as specified in [RFC 9728](https://datatracker.ietf.org/doc/html/rfc9728). This endpoint provides +clients with information about the resource server's OAuth configuration, including which +authorization servers can be used and what scopes are supported. + +The SDK provides [`ProtectedResourceMetadataHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#ProtectedResourceMetadataHandler) +to serve this metadata. The handler automatically sets CORS headers (`Access-Control-Allow-Origin: *`) +to support cross-origin client discovery, as the metadata contains only public configuration information. + +Example usage: + +```go +metadata := &oauthex.ProtectedResourceMetadata{ + Resource: "https://example.com/mcp", + AuthorizationServers: []string{ + "https://auth.example.com/.well-known/openid-configuration", + }, + ScopesSupported: []string{"read", "write"}, +} +http.Handle("/.well-known/oauth-protected-resource", + auth.ProtectedResourceMetadataHandler(metadata)) +``` + +For more sophisticated CORS policies, wrap the handler with a CORS middleware like +[github.com/rs/cors](https://github.com/rs/cors) or [github.com/jub0bs/cors](https://github.com/jub0bs/cors). The [_auth middleware example_](https://github.com/modelcontextprotocol/go-sdk/tree/main/examples/server/auth-middleware) shows how to implement authorization for both JWT tokens and API keys. ### Client -Client-side OAuth is implemented by setting -[`StreamableClientTransport.HTTPClient`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk@v0.5.0/mcp#StreamableClientTransport.HTTPClient) to a custom [`http.Client`](https://pkg.go.dev/net/http#Client) -Additional support is forthcoming; see modelcontextprotocol/go-sdk#493. +Client-side authorization is supported via the +[`StreamableClientTransport.OAuthHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#StreamableClientTransport.OAuthHandler) +field. If the handler is provided, the transport will automatically use it to +add an `Authorization: Bearer ` header to every request. The transport +will also call the handler's `Authorize` method if the server returns +`401 Unauthorized` or `403 Forbidden` errors to perform the authorization flow +or facilitate scope step-up authorization. + +The SDK implements the Authorization Code flow in +[`auth.AuthorizationCodeHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#AuthorizationCodeHandler). +This handler supports: + +- [Client ID Metadata Documents](https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#client-id-metadata-documents) +- [Pre-registered clients](https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#preregistration) +- [Dynamic Client Registration](https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#dynamic-client-registration) +- [RFC 9207](https://www.rfc-editor.org/rfc/rfc9207) Authorization Server Issuer Identification + +To use it, configure the handler and assign it to the transport: + +```go +authHandler, _ := auth.NewAuthorizationCodeHandler(&auth.AuthorizationCodeHandlerConfig{ + RedirectURL: "https://myapp.com/oauth2-callback", + // Configure one of the following: + // ClientIDMetadataDocumentConfig: ... + // PreregisteredClientConfig: ... + // DynamicClientRegistrationConfig: ... + AuthorizationCodeFetcher: func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + // Open the args.URL in a browser and return the resulting code, state, and iss. + // See full example in examples/auth/client/main.go. + code := ... + state := ... + iss := ... // "iss" query parameter from the redirect URI (RFC 9207) + return &auth.AuthorizationResult{Code: code, State: state, Iss: iss}, nil + }, +}) + +transport := &mcp.StreamableClientTransport{ + Endpoint: "https://example.com/mcp", + OAuthHandler: authHandler, +} +client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) +session, err := client.Connect(ctx, transport, nil) +``` + +The `auth.AuthorizationCodeHandler` automatically manages token refreshing (if the server provides a refresh token) and step-up authentication (when the server returns `insufficient_scope` error). + +#### Enterprise Managed Authorization (SEP-990) + +For enterprise SSO scenarios where users authenticate with an enterprise Identity Provider (IdP), +the SDK provides +[`extauth.EnterpriseHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth/extauth#EnterpriseHandler), +an implementation of `OAuthHandler` that automates the Enterprise Managed Authorization flow: + +1. **OIDC Login**: User authenticates with enterprise IdP → ID Token +2. **Token Exchange** (RFC 8693): ID Token → ID-JAG at IdP +3. **JWT Bearer Grant** (RFC 7523): ID-JAG → Access Token at MCP Server + +To use enterprise managed authorization, create an `EnterpriseHandler` and assign it to your transport: + +```go +// Create ID token fetcher using OIDC login +idTokenFetcher := func(ctx context.Context) (*oauth2.Token, error) { + oidcConfig := &extauth.OIDCLoginConfig{ + IssuerURL: "https://company.okta.com", + Credentials: &oauthex.ClientCredentials{ + ClientID: "idp-client-id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "idp-client-secret", + }, + }, + RedirectURL: "http://localhost:3142", + Scopes: []string{"openid", "profile", "email"}, + } + + tokens, err := extauth.PerformOIDCLogin(ctx, oidcConfig, authCodeFetcher) + if err != nil { + return nil, err + } + + return tokens, nil +} + +// Create Enterprise Handler +enterpriseHandler, err := extauth.NewEnterpriseHandler(&extauth.EnterpriseHandlerConfig{ + IdPIssuerURL: "https://company.okta.com", + IdPCredentials: &oauthex.ClientCredentials{ + ClientID: "idp-client-id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "idp-client-secret", + }, + }, + MCPAuthServerURL: "https://auth.mcpserver.example", + MCPResourceURI: "https://mcp.mcpserver.example", + MCPCredentials: &oauthex.ClientCredentials{ + ClientID: "mcp-client-id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ + ClientSecret: "mcp-client-secret", + }, + }, + MCPScopes: []string{"read", "write"}, + IDTokenFetcher: idTokenFetcher, +}) + +// Use with transport +transport := &mcp.StreamableClientTransport{ + Endpoint: "https://example.com/mcp", + OAuthHandler: enterpriseHandler, +} +client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) +session, err := client.Connect(ctx, transport, nil) +``` + +The `EnterpriseHandler` automatically manages the token exchange flow. Note that it intentionally does not support refresh tokens - when an access token expires, the entire authorization flow is repeated to ensure enterprise policies are consistently enforced. + +For a complete working example, see [examples/auth/enterprise](https://github.com/modelcontextprotocol/go-sdk/tree/main/examples/auth/enterprise). ## Security Here we discuss the mitigations described under -the MCP spec's [Security Best Practices](https://modelcontextprotocol.io/specification/2025-06-18/basic/security_best_practices) section, and how we handle them. +the MCP's [Security Best Practices](https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices) section, and how we handle them. ### Confused Deputy -The [mitigation](https://modelcontextprotocol.io/specification/2025-06-18/basic/security_best_practices#mitigation), obtaining user consent for dynamically registered clients, -happens on the MCP client. At present we don't provide client-side OAuth support. - +The [mitigation](https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#mitigation), +obtaining user consent for dynamically registered clients, is mostly the +responsibility of the MCP Proxy server implementation. The SDK client does +generate cryptographically secure random `state` values for each authorization +request by default and validates them when the authorization code is returned. +Mismatched state values will result in an error. ### Token Passthrough -The [mitigation](https://modelcontextprotocol.io/specification/2025-06-18/basic/security_best_practices#mitigation-2), accepting only tokens that were issued for the server, depends on the structure +The [mitigation](https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#mitigation-2), accepting only tokens that were issued for the server, depends on the structure of tokens and is the responsibility of the [`TokenVerifier`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#TokenVerifier) provided to [`RequireBearerToken`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#RequireBearerToken). +### Server-Side Request Forgery + +The OAuth discovery helpers apply these [mitigations](https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#mitigation-3) by default: + +- _Enforce HTTPS_. Reject `http://` URLs (except loopback: `localhost`, `127.0.0.1`, `::1`) +and redirects that downgrade `https` to a non-`https` scheme. + +- _Block Private IP Ranges_. Reject targets that resolve to a non-public address +(private, link-local including the `169.254.169.254` metadata endpoint, CGNAT, multicast, +unspecified). The check runs on the initial URL and at dial time, guarding against DNS rebinding. + +- _Validate Redirect Targets_. Reject redirects to non-public addresses and cap the redirect count. + +- _Use Egress Proxies_. Out of scope for the SDK; configure separately. + +**Opting out.** Passing a custom `http.Client` can bypass these checks, making SSRF +protection your responsibility. Specifically: a custom `Transport.DialContext`/`DialTLSContext`, +a non-`*http.Transport`, or a configured `Proxy` disables the dial-time IP check; setting +`CheckRedirect` replaces all redirect validation. In these cases, implement your own IP +blocking, redirect validation, and DNS pinning. + ### Session Hijacking -The [mitigations](https://modelcontextprotocol.io/specification/2025-06-18/basic/security_best_practices#mitigation-3) are as follows: +The [mitigations](https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#mitigation-4) are as follows: - _Verify all inbound requests_. The [`RequireBearerToken`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#RequireBearerToken) middleware function will verify all HTTP requests that it receives. It is the @@ -226,12 +564,36 @@ user's responsibility to wrap that function around all handlers in their server. - _Secure session IDs_. This SDK generates cryptographically secure session IDs by default. If you create your own with [`ServerOptions.GetSessionID`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.GetSessionID), it is your responsibility to ensure they are secure. -If you are using Go 1.24 or above, -we recommend using [`crypto/rand.Text`](https://pkg.go.dev/crypto/rand#Text) +We recommend using [`crypto/rand.Text`](https://pkg.go.dev/crypto/rand#Text). -- _Binding session IDs to user information_. This is an application requirement, out of scope -for the SDK. You can create your own session IDs by setting -[`ServerOptions.GetSessionID`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.GetSessionID). +- _Binding session IDs to user information_. The SDK supports this mitigation through +[`TokenInfo.UserID`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#TokenInfo.UserID). +When a [`TokenVerifier`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#TokenVerifier) +sets `UserID` on the returned `TokenInfo`, the streamable transport will: + 1. Store the user ID when a new session is created. + 2. Verify that subsequent requests to that session include a token with the same `UserID`. + 3. Reject requests with a 403 Forbidden if the user ID doesn't match. + + **Recommendation**: If your `TokenVerifier` can extract a user identifier from the token + (such as a `sub` claim in a JWT, or a user ID associated with an API key), set + `TokenInfo.UserID` to enable this protection. This prevents an attacker with a valid + token from hijacking another user's session by guessing or obtaining their session ID. + +### Issuer Mix-Up + +The [mitigation](https://www.rfc-editor.org/rfc/rfc9207) against issuer mix-up attacks is +implemented per [RFC 9207](https://www.rfc-editor.org/rfc/rfc9207). The SDK client validates +the `iss` parameter in authorization responses to ensure they originated from the expected +authorization server: + +- If `iss` is present in the redirect URI, the SDK verifies it matches the issuer from the + authorization server's metadata. A mismatch results in an error. +- If `iss` is absent but the authorization server advertises + `authorization_response_iss_parameter_supported: true` in its [RFC 8414](https://www.rfc-editor.org/rfc/rfc8414) + metadata, the SDK rejects the response with an error. + +The `AuthorizationCodeFetcher` is responsible for extracting the `iss` query parameter from +the redirect URI and returning it in [`AuthorizationResult.Iss`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/auth#AuthorizationResult). ## Utilities @@ -263,6 +625,13 @@ session if the ping fails, set or [`ServerOptions.KeepAlive`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.KeepAlive). +> **Note**: `ping` is removed from the protocol as of `2026-07-28` by +> [SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575). +> The SDK preserves the API for backward compatibility with legacy clients; +> when both peers negotiate `2026-07-28`, the server returns +> `MethodNotFound` (`-32601`) for `ping` and `KeepAlive` should not be +> enabled. + ### Progress [Progress](https://modelcontextprotocol.io/specification/2025-06-18/basic/utilities/progress) @@ -279,3 +648,23 @@ or Issue #460 discusses some potential ergonomic improvements to this API. %include ../../mcp/mcp_example_test.go progress - + +### Error codes + +The SDK uses the standard JSON-RPC base codes (parse error `-32700`, +invalid request `-32600`, method not found `-32601`, invalid params +`-32602`, internal error `-32603`) plus the following MCP-specific codes: + +| Constant | Code | Meaning | +|---|---|---| +| [`CodeHeaderMismatch`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#CodeHeaderMismatch) | `-32020` | An MCP HTTP header (`Mcp-Method`, `Mcp-Name`, `Mcp-Protocol-Version`, or `Mcp-Param-*`) does not match the JSON-RPC body | +| [`CodeMissingRequiredClientCapabilities`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#CodeMissingRequiredClientCapabilities) | `-32021` | Server requires client capabilities the client did not declare | +| [`CodeUnsupportedProtocolVersion`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#CodeUnsupportedProtocolVersion) | `-32022` | Requested protocol version not supported. Data: `UnsupportedProtocolVersionData{Supported []string, Requested string}` | +| [`CodeResourceNotFound`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#CodeResourceNotFound) | `-32602` | Resource URI not found (variable; was `-32002` before SEP-2164) | + +The error code allocation policy defined in `2026-07-28` +([SEP-2575](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2575)) +partitions the JSON-RPC server-error range: + +- `-32000` to `-32019`: implementation-defined; existing SDK usage is grandfathered. +- `-32020` to `-32099`: reserved for the MCP specification. diff --git a/internal/docs/quick_start.src.md b/internal/docs/quick_start.src.md new file mode 100644 index 0000000..c0f1546 --- /dev/null +++ b/internal/docs/quick_start.src.md @@ -0,0 +1,27 @@ +# Quick Start + +%toc + +## Installation + +``` +go get github.com/modelcontextprotocol/go-sdk/mcp +``` + +## Getting started + +To get started creating an MCP server, create an `mcp.Server` instance, add +features to it, and then run it over an `mcp.Transport`. For example, this +server adds a single simple tool, and then connects it to clients over +stdin/stdout: + +%include ../readme/server/server.go - + +To communicate with that server, create an `mcp.Client` and connect it to the +corresponding server, by running the server command and communicating over its +stdin/stdout: + +%include ../readme/client/client.go - + +The [`examples/`](https://github.com/modelcontextprotocol/go-sdk/tree/main/examples) directory contains more example clients and +servers. diff --git a/internal/docs/rough_edges.src.md b/internal/docs/rough_edges.src.md new file mode 100644 index 0000000..90751f7 --- /dev/null +++ b/internal/docs/rough_edges.src.md @@ -0,0 +1,69 @@ +# Rough edges: API decisions to reconsider for v2 + +This file collects a list of API oversights or rough edges that we've uncovered +post v1.0.0, along with their current workarounds. These issues can't be +addressed without breaking backward compatibility, but we'll revisit them for +v2. + +- `EventStore.Open` is unnecessary. This was an artifact of an earlier version + of the SDK where event persistence and delivery were combined. + + **Workaround**: `Open` may be implemented as a no-op. + +- `Event` need not have been exported: it's an implementation detail of the SSE + and streamable transports. Also the 'Name' field is a misnomer: it should be + 'event'. + +- Enforcing valid tool names: with + [SEP-986](https://github.com/modelcontextprotocol/modelcontextprotocol/issues/986) + landing after the SDK was at v1, we missed an opportunity to panic on invalid + tool names. Instead, we have to simply produce an error log. In v2, we should + panic. + +- Inconsistent naming. + - `ResourceUpdatedNotificationsParams` should probably have just been + `ResourceUpdatedParams`, as we don't include the word 'notification' in + other notification param types. + - Similarly, `ProgressNotificationParams` should probably have been + `ProgressParams`. + +- `AudioContent.MarshalJSON` should have had a pointer receiver, to be + consistent with other content types. + +- `ClientCapabilities.Roots` should have been a distinguished struct pointer + ([see #607](https://github.com/modelcontextprotocol/go-sdk/issues/607)). + + **Workaround**: use `ClientCapabilities.RootsV2`, which aligns with the + semantics of other capability fields. + +- Default capabilities should have been empty. Instead, servers default to + advertising `logging`, and clients default to advertising `roots` with + `listChanged: true`. This is confusing because a nil `Capabilities` field + does not mean "no capabilities". + + **Workaround**: to advertise no capabilities, set + `ServerOptions.Capabilities` or `ClientOptions.Capabilities` to an empty + `&ServerCapabilities{}` or `&ClientCapabilities{}` respectively. + +- `CreateMessageResult.Content` is singular `Content`, but the 2025-11-25 spec + allows `content` to be a single block or an array (for parallel tool calls). + We added `CreateMessageResultWithTools` (with `Content []Content`) as a + workaround, matching the TypeScript SDK's approach. In v2, + `CreateMessageResult` should use `[]Content` directly. Similarly, + `SamplingMessage.Content` should become `[]Content` to support sending + multiple tool_result blocks in a single user message. + +- We didn't actually need CallToolParams and CallToolParamsRaw, since even when + we're unmarshalling into a custom Go type (for the mcp.AddTool convenience + wrapper) we need to first unmarshal into a `map[string]any` in order to do + server-side validation of required fields. CallToolParams could have just had + a map[string]any. + +- `StreamableHTTPOptions.CrossOriginProtection` should not have been part of + the SDK API. Cross-origin protection is a general HTTP concern, not specific + to MCP, and can be applied as standard HTTP middleware. + +- `ToolAnnotations` (`mcp/protocol.go`) should have all fields typed as `*bool` + for full control to define what is being sent over the wire. Different + MCP clients have different requirements, and some of them require all fields + to be explicitly set to either `true` or `false`. \ No newline at end of file diff --git a/internal/docs/server.src.md b/internal/docs/server.src.md index c09ba63..ff66383 100644 --- a/internal/docs/server.src.md +++ b/internal/docs/server.src.md @@ -33,6 +33,24 @@ notification. %include ../../mcp/server_example_test.go prompts - +### Prompt message content + +A +[`PromptMessage`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#PromptMessage) +has two fields: a `Role`, either `"user"` or `"assistant"`, and a single +`Content`. That is the same interface a tool result uses (see [Tool result +content](#tool-result-content)), so a prompt message can hold an image, audio, +or a resource's contents as readily as text. + +Media on its own says nothing about what the model should do with it, so send a +text message alongside it. Embedding a resource puts its contents directly in +the message, sparing the client a separate `resources/read`. Note that +`ResourceContents.URI` is not checked against the resources the server has +registered: it records where the contents came from, so use the URI of a real +resource when there is one, as in the example below. + +%include ../../mcp/server_example_test.go promptcontent - + ## Resources In MCP terms, a _resource_ is some data referenced by a URI. @@ -65,7 +83,12 @@ and to unsubscribe. Set [`ClientOptions.ResourceUpdatedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.ResourceUpdatedHandler) -to be notified of changes to subscribed resources. +to be notified of changes to subscribed resources. On `2026-07-28` and +later sessions the SDK delivers these notifications over a +`subscriptions/listen` stream instead of the legacy `resources/subscribe` +RPC; see [Subscriptions +(`subscriptions/listen`)](protocol.md#subscriptions) +for the wire-level details. **Server-side**: Use @@ -89,6 +112,37 @@ notification. %include ../../mcp/server_example_test.go resources - +### Binary resources + +A +[`ResourceContents`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ResourceContents) +carries its data in either `Text` or `Blob`. Binary data goes in `Blob` as +plain bytes — the SDK base64-encodes it on the wire — and leaves `Text` empty, +so a client tells the two apart by which field is populated. + +%include ../../mcp/server_example_test.go binaryresource - + +### Resource subscriptions + +A client interested in changes to one resource subscribes to its URI with +[`ClientSession.Subscribe`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientSession.Subscribe) +and stops with +[`ClientSession.Unsubscribe`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientSession.Unsubscribe). +Set +[`ServerOptions.SubscribeHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.SubscribeHandler) +and +[`UnsubscribeHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.UnsubscribeHandler) +to track who is listening; setting either one gives the server the +`resources.subscribe` capability. + +Publish a change with +[`Server.ResourceUpdated`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server.ResourceUpdated), +which notifies every session subscribed to that URI. The notification carries +only the URI, not the new contents, so a client that wants them reads the +resource again. + +%include ../../mcp/server_example_test.go subscribe - + ## Tools MCP servers can provide @@ -140,7 +194,7 @@ In order to implement a tool, the user must do all of the following: - Unmarshal the input schema into a Go value - Execute the tool logic. - Marshal the tool's structured output (if any) to JSON, and store it in the - result's `StructuredOutput` field as well as the unstructured `Content` field. + result's `StructuredContent` field as well as the unstructured `Content` field. - Validate that output JSON against the tool's output schema. - If any tool errors occurred, pack them into the unstructured content and set `IsError` to `true.` @@ -165,13 +219,17 @@ mcp.AddTool(server, &mcp.Tool{Name: "my_tool"}, handler) This does the following automatically: -- If `Tool.InputSchema` or `Tool.OutputSchema` are unset, the input and output - schemas are inferred from the `In` type, which must be a struct or map. - Optional `jsonschema` struct tags provide argument descriptions. +- If `Tool.InputSchema` is unset, the input schema is inferred from the `In` + type, which must be a struct or map. +- If `Tool.OutputSchema` is unset and the `Out` type is not `any`, the output + schema is inferred from the `Out` type. Per SEP-2106, `Out` may be any Go + type whose inferred schema is a valid JSON Schema (struct, map, slice, + primitive, etc.). +- Optional `jsonschema` struct tags provide argument and output descriptions. - Tool arguments are validated against the input schema. - Tool arguments are marshaled into the `In` value. - Tool output (the `Out` value) is marshaled into the result's - `StructuredOutput`, as well as the unstructured `Content`. + `StructuredContent`, as well as the unstructured `Content`. - Output is validated against the tool's output schema. - If an ordinary error is returned, it is stored int the `CallToolResult` and `IsError` is set to `true`. @@ -189,10 +247,160 @@ the `Tool.InputSchema` explicitly: %include ../../mcp/tool_example_test.go customschemas - -_See [mcp/tool_example_test.go](../mcp/tool_example_test.go) for the full -example, or [examples/server/toolschemas](examples/server/toolschemas/main.go) +_See [mcp/tool_example_test.go](https://github.com/modelcontextprotocol/go-sdk/blob/main/mcp/tool_example_test.go) for the full +example, or [examples/server/toolschemas](https://github.com/modelcontextprotocol/go-sdk/blob/main/examples/server/toolschemas/main.go) for more examples of customizing tool schemas._ +### Tool result content + +Alongside its structured output, a tool result carries a list of +[`Content`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Content) +blocks in `CallToolResult.Content`, and may mix as many kinds as it needs: + +| Type | Carries | +| --- | --- | +| [`TextContent`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#TextContent) | plain text | +| [`ImageContent`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ImageContent) | image data, with a MIME type | +| [`AudioContent`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#AudioContent) | audio data, with a MIME type | +| [`EmbeddedResource`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#EmbeddedResource) | the contents of a resource, inline | +| [`ResourceLink`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ResourceLink) | a reference to a resource the client can read separately | + +`ImageContent.Data` and `AudioContent.Data` are plain `[]byte`; the SDK +base64-encodes them on the wire, so handlers never encode by hand, and base64 +carried over from another MCP SDK must be decoded before it is assigned. An +`EmbeddedResource` wraps a +[`ResourceContents`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ResourceContents), +which holds either `Text` or, for binary data, `Blob`. Prefer a `ResourceLink` +when the client may not need the contents, since an embedded resource is +transferred whether it is used or not. + +When a handler bound with `AddTool` leaves `Content` unset, the SDK fills it +with the JSON encoding of the output value; set `Content` explicitly, as below, +to return anything else. `StructuredContent` is still populated from the output +value either way. + +%include ../../mcp/tool_example_test.go contenttypes - + +`ToolUseContent` and `ToolResultContent` also implement `Content`, but are only +valid in sampling messages, not in tool results. + +**Stateless server deployments:** Some deployments create a new +[`Server`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server) +for each incoming request, re-registering tools every time. To avoid repeated +schema generation, create a +[`SchemaCache`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#SchemaCache) +and share it across server instances: + +```go +var schemaCache = mcp.NewSchemaCache() // create once at startup + +func handleRequest(w http.ResponseWriter, r *http.Request) { + s := mcp.NewServer(impl, &mcp.ServerOptions{SchemaCache: schemaCache}) + mcp.AddTool(s, myTool, myHandler) + // ... +} +``` + +## List changed notifications + +Adding or removing a feature on a connected server sends the matching +`notifications/*/list_changed` to every client, which dispatches it to +[`ClientOptions.ToolListChangedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.ToolListChangedHandler), +[`PromptListChangedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.PromptListChangedHandler), +or +[`ResourceListChangedHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.ResourceListChangedHandler). +A notification reports only that the list changed, so a client that needs the +new contents lists them again. + +These notifications ride on capabilities, and capabilities are inferred from +what is registered before [`Server.Connect`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server.Connect) +(see [Capabilities](#capabilities)). A server that registers its features +afterwards, as the example below does, has to declare +`ServerOptions.HasTools`, `HasPrompts`, or `HasResources` itself, or the +client is never told. + +%include ../../mcp/server_example_test.go listchanged - + +## Multi Round-Trip Requests + +[SEP-2322](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2322) +defines a new pattern for server-to-client requests (sampling, elicitation, +roots). Instead of issuing a fresh JSON-RPC request mid-flight, the server +returns an `InputRequiredResult` from its in-flight handler — the +`InputRequests` field of `CallToolResult`, `GetPromptResult`, or +`ReadResourceResult` carries the requests for additional information. The +client responds by retrying the original call with `InputResponses` set. + +The SDK supports this pattern from both sides without requiring callers to +choose: + +- **Returning input requests from a handler.** Set `InputRequests` on the + result you return. + Leave `Content` / `StructuredContent` empty — returning + both at once is a server bug and the SDK returns `-32603 InternalError`. +- **Calling the legacy APIs.** `ServerSession.Elicit`, + `ServerSession.CreateMessage(WithTools)`, and `ServerSession.ListRoots` + remain available and work for both new and legacy clients. + +### Talking to legacy clients + +The server installs `serverMultiRoundTripMiddleware` automatically. For +clients on a protocol version earlier than `2026-07-28`, the middleware +intercepts any `InputRequiredResult` your handler returns, fulfils each +input request itself by calling the legacy server-initiated APIs +(`Elicit`, `CreateMessage`, `ListRoots`), and re-invokes your handler +exactly once with the responses already populated. This means a handler +written in the MRTR style works against both old and new clients without +code changes. + +### Example + +The following example shows a "greet" tool whose handler asks the user +for their name via elicitation before producing the final greeting. The +handler runs twice — once to issue the elicitation, once to consume the +response — but the client call site sees a single `CallTool` returning the +final result, because the SDK's MRTR middleware handles the round trip on +either side of the wire (depending on the negotiated protocol version). + +%include ../../mcp/server_example_test.go mrtr - + +## Cacheable list results + +[SEP-2549](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2549) +adds `ttlMs` and `cacheScope` fields to the results of `tools/list`, +`prompts/list`, `resources/list`, `resources/templates/list`, and +`resources/read`. Both fields complement existing list +results: clients use `ttlMs` as a freshness hint to reduce polling, +and `cacheScope` (`"public"` or `"private"`) controls whether shared +intermediaries may cache the response. + +**Server-side defaults.** The SDK fills in `CacheScope = "public"` whenever a +result leaves it empty, since the field is required on the wire. Everything +else is left as produced: a `ResourceHandler` that sets `TTLMs` or `CacheScope` +on its `ReadResourceResult` keeps those values. The list results and +`server/discover` are built by the SDK itself, so they start out with +`TTLMs = 0`, which a client cache treats as "immediately stale". + +**Setting a policy.** `ServerOptions.SetCacheable` decides the fields for every +result that carries them. It runs after the corresponding handler, receives +what that handler produced, and may keep, adjust, or replace it. The request is +passed in so that policy can vary by method, session, or authorization context. + +```go +server := mcp.NewServer(impl, &mcp.ServerOptions{ + SetCacheable: func(_ context.Context, req mcp.Request, c *mcp.Cacheable) { + // A 30s freshness hint, except where the handler chose its own. + if c.TTLMs == 0 { + c.TTLMs = 30_000 + } + }, +}) +``` + +`SetCacheable` may run with the server's lock held, so it must not call back +into the `Server`: adding or removing a feature, or ranging over +`Server.Sessions`, deadlocks. + ## Utilities ### Completion @@ -215,6 +423,13 @@ requests. ### Logging +> **Note:** The logging feature is deprecated as of protocol version +> 2026-07-28 +> ([SEP-2577](https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging)). +> It remains fully functional during the deprecation window (at least twelve +> months). The SDK continues to support logging for compatibility. Servers +> should migrate to stderr logging (for STDIO transports) or OpenTelemetry. + MCP servers can send logging messages to MCP clients. (This form of logging is distinct from server-side logging, where the server produces logs that remain server-side, for use by server maintainers.) @@ -222,8 +437,12 @@ server produces logs that remain server-side, for use by server maintainers.) **Server-side**: The minimum log level is part of the server state. For stateful sessions, there is no default log level: no log messages will be sent -until the client calls `SetLevel` (see below). -For stateful sessions, the level defaults to "info". +until the client calls `SetLoggingLevel` (see below). +For legacy stateless sessions the level defaults to "info". +For sessionless 2026-07-28 requests there is no server-side default: each request +carries its level in `_meta` (`io.modelcontextprotocol/logLevel`), and nothing is +logged until the client sends one. +(Note: logging is deprecated as of 2026-07-28 / SEP-2577.) [`ServerSession.Log`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerSession.Log) is the low-level way for servers to log to clients. It sends a logging notification to the client if the level of the message @@ -239,10 +458,65 @@ Servers always report the logging capability. **Client-side**: Set [`ClientOptions.LoggingMessageHandler`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientOptions.LoggingMessageHandler) to receive log messages. -Call [`ClientSession.SetLevel`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientSession.SetLevel) to change the log level for a session. +Call [`ClientSession.SetLoggingLevel`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ClientSession.SetLoggingLevel) to change the log level for a session. %include ../../mcp/server_example_test.go logging - +## Capabilities + +Server capabilities are advertised to clients during the initialization +handshake. By default, the SDK advertises only the `logging` capability. +Additional capabilities are automatically added when features are registered +(e.g., adding a tool adds the `tools` capability). + +### Capability inference + +When features such as tools, prompts, or resources are added to the server +(e.g., via `Server.AddTool`), their capability is automatically inferred, with +default value `{listChanged:true}`. Similarly, if the +`ServerOptions.SubscribeHandler` or `ServerOptions.CompletionHandler` are set, +the corresponding capability is added. + +### Explicit capabilities + +To explicitly declare capabilities, or to override the [default inferred +capability](#capability-inference), set +[`ServerOptions.Capabilities`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#ServerOptions.Capabilities). +This sets the default server capabilities, before any capabilities are added +based on configured handlers. If a capability is already present as a field in +`Capabilities`, adding a feature or handler will not change its configuration. + +This allows you to: + +- **Disable default capabilities**: Pass an empty `&ServerCapabilities{}` to + disable all defaults, including logging. +- **Disable listChanged notifications**: Set `ListChanged: false` on a + capability to prevent the server from sending list-changed notifications + when features are added or removed. +- **Pre-declare capabilities**: Declare capabilities before features are + registered, useful for servers that load features dynamically. + +```go +// Disable listChanged notifications for tools +server := mcp.NewServer(impl, &mcp.ServerOptions{ + Capabilities: &mcp.ServerCapabilities{ + Logging: &mcp.LoggingCapabilities{}, + Tools: &mcp.ToolCapabilities{ListChanged: false}, + }, +}) +``` + +**Deprecated**: The `HasPrompts`, `HasResources`, and `HasTools` fields on +`ServerOptions` are deprecated. Use `Capabilities` instead. + +### Extensions + +[SEP-2133](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2133) +adds an `extensions` map to `ServerCapabilities` so that optional +capabilities outside the core protocol can be declared on the wire. Keys +are namespaced as `"{vendor-prefix}/{extension-name}"`; values are +per-extension settings objects. + ### Pagination Server-side feature lists may be diff --git a/internal/docs/toc_test.go b/internal/docs/toc_test.go new file mode 100644 index 0000000..23b02e7 --- /dev/null +++ b/internal/docs/toc_test.go @@ -0,0 +1,70 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package docs + +import ( + "os" + "path/filepath" + "regexp" + "strings" + "testing" +) + +// Headings are slugged twice by two different tools: weave writes the %toc +// links, and the renderer (GitHub, and mkdocs on go.sdk.modelcontextprotocol.io) +// assigns the heading ids. The two agree on plain words, and disagree on +// punctuation: weave keeps "(", ")" and "/" and drops "_", the renderers drop +// the first three and keep the last. A heading carrying any of them therefore +// gets a table-of-contents entry that lands nowhere. +// +// This test reads the generated docs and checks every %toc entry against the +// ids the renderers will actually produce. + +var ( + headingRE = regexp.MustCompile(`(?m)^(#{1,6})\s+(.*?)\s*$`) + tocLineRE = regexp.MustCompile(`(?m)^\s*1\. \[(.*?)\]\((#.*?)\)\s*$`) + dropRE = regexp.MustCompile(`[^\w\s-]`) + spaceRE = regexp.MustCompile(`[-\s]+`) +) + +// renderedID reproduces the anchor GitHub and python-markdown's toc extension +// derive from heading text: take the text, drop everything that is not a word +// character, whitespace or a hyphen, lowercase it, and join on hyphens. +func renderedID(heading string) string { + s := strings.ReplaceAll(heading, "`", "") + s = dropRE.ReplaceAllString(s, "") + s = strings.ToLower(strings.TrimSpace(s)) + return spaceRE.ReplaceAllString(s, "-") +} + +func TestTOCAnchorsResolve(t *testing.T) { + files, err := filepath.Glob(filepath.Join("..", "..", "docs", "*.md")) + if err != nil { + t.Fatal(err) + } + if len(files) == 0 { + t.Fatal("no generated docs found") + } + for _, file := range files { + data, err := os.ReadFile(file) + if err != nil { + t.Fatal(err) + } + text := string(data) + + ids := make(map[string]bool) + for _, m := range headingRE.FindAllStringSubmatch(text, -1) { + ids[renderedID(m[2])] = true + } + for _, m := range tocLineRE.FindAllStringSubmatch(text, -1) { + label, target := m[1], strings.TrimPrefix(m[2], "#") + if !ids[target] { + t.Errorf("%s: table-of-contents entry %q points at #%s, which no heading produces "+ + "(want #%s); drop (, ), / and _ from the heading", + filepath.Base(file), label, target, renderedID(label)) + } + } + } +} diff --git a/internal/json/json.go b/internal/json/json.go new file mode 100644 index 0000000..3a0e52b --- /dev/null +++ b/internal/json/json.go @@ -0,0 +1,93 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +// Package json provides internal JSON utilities. + +package json + +import ( + "bytes" + "fmt" + "io" + + "github.com/segmentio/encoding/json" +) + +// defaultMaxDepth is the maximum JSON nesting depth accepted by [Unmarshal]. +// Realistic MCP payloads nest only a handful of levels, so 1000 is far more +// than any legitimate message requires while keeping worst-case parse cost +// bounded. +const defaultMaxDepth = 1000 + +// errMaxDepthExceeded is returned by [Unmarshal] when the input nests deeper than [defaultMaxDepth]. +var errMaxDepthExceeded = fmt.Errorf("json: exceeded maximum nesting depth of %d", defaultMaxDepth) + +type Decoder struct { + dec *json.Decoder +} + +func NewDecoder(r io.Reader) *Decoder { + dec := json.NewDecoder(r) + dec.DontMatchCaseInsensitiveStructFields() + return &Decoder{dec: dec} +} + +func (d *Decoder) Decode(v any) error { + return d.dec.Decode(v) +} + +func Unmarshal(data []byte, v any) error { + if err := checkMaxDepth(data, defaultMaxDepth); err != nil { + return err + } + return NewDecoder(bytes.NewReader(data)).Decode(v) +} + +func UnmarshalUseNumber(data []byte, v any) error { + if err := checkMaxDepth(data, defaultMaxDepth); err != nil { + return err + } + dec := NewDecoder(bytes.NewReader(data)) + dec.dec.UseNumber() + return dec.Decode(v) +} + +// checkMaxDepth scans data once and reports [errMaxDepthExceeded] if the +// nesting of JSON objects and arrays exceeds maxDepth. It is a lightweight pass which +// tracks '{' and '[' against '}' and ']' while skipping over the contents of strings. +func checkMaxDepth(data []byte, maxDepth int) error { + var ( + depth int + inString bool + escaped bool + ) + for _, b := range data { + if inString { + switch { + case escaped: + escaped = false + case b == '\\': + escaped = true + case b == '"': + inString = false + } + continue + } + + switch b { + case '"': + inString = true + case '{', '[': + depth++ + if depth > maxDepth { + return errMaxDepthExceeded + } + case '}', ']': + if depth > 0 { + depth-- + } + } + } + return nil +} diff --git a/internal/json/json_test.go b/internal/json/json_test.go new file mode 100644 index 0000000..b1d110c --- /dev/null +++ b/internal/json/json_test.go @@ -0,0 +1,259 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package json + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "strings" + "testing" + "time" + + "github.com/google/go-cmp/cmp" +) + +func TestUnmarshalCaseSensitivity(t *testing.T) { + type Nested struct { + Field string `json:"field"` + } + type Target struct { + Field string + TaggedField string `json:"custom_tag"` + Nested *Nested + } + + tests := []struct { + name string + json string + want Target + }{ + { + name: "exact match", + json: `{"Field": "value", "custom_tag": "tagged", "Nested": {"field": "nested"}}`, + want: Target{ + Field: "value", + TaggedField: "tagged", + Nested: &Nested{ + Field: "nested", + }, + }, + }, + { + name: "case mismatch", + json: `{"field": "value", "Custom_tag": "tagged", "Nested": {"Field": "nested"}}`, + want: Target{ + Field: "", + TaggedField: "", + Nested: &Nested{ + Field: "", + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got Target + if err := Unmarshal([]byte(tt.json), &got); err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + if diff := cmp.Diff(tt.want, got); diff != "" { + t.Errorf("Unmarshal mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestNewDecoderCaseSensitivity(t *testing.T) { + type Target struct { + Field string `json:"field"` + TaggedField string `json:"custom_tag"` + } + + tests := []struct { + name string + json string + want Target + }{ + { + name: "exact match", + json: `{"field": "value", "custom_tag": "tagged"}`, + want: Target{Field: "value", TaggedField: "tagged"}, + }, + { + name: "case mismatch", + json: `{"Field": "value", "Custom_tag": "tagged"}`, + want: Target{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got Target + dec := NewDecoder(strings.NewReader(tt.json)) + if err := dec.Decode(&got); err != nil { + t.Fatalf("Decode failed: %v", err) + } + if diff := cmp.Diff(tt.want, got); diff != "" { + t.Errorf("Decode mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestUnmarshalNullCharacter(t *testing.T) { + type Target struct { + Field string `json:"field"` + } + + tests := []struct { + name string + json string + want string + }{ + { + name: "null char in middle", + json: `{"fi\u0000eld": "value"}`, + }, + { + name: "null char at end", + json: `{"field\u0000": "value"}`, + }, + { + name: "null char at start", + json: `{"\u0000field": "value"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got Target + if err := Unmarshal([]byte(tt.json), &got); err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + if got.Field != "" { + t.Errorf("Field was set: %s", tt.json) + } + }) + } +} + +func TestUnmarshalChecksNestingDepth(t *testing.T) { + createNesting := func(depth int) []byte { + return []byte(strings.Repeat("[", depth) + strings.Repeat("]", depth)) + } + + tests := []struct { + name string + data []byte + wantErr bool + }{ + {"flat object", []byte(`{"a":1}`), false}, + {"at limit", createNesting(defaultMaxDepth), false}, + {"one over limit", createNesting(defaultMaxDepth + 1), true}, + {"brackets in strings do not count", []byte(`{"a":"[[[[[[[[[["}`), false}, + {"escaped quote in string", []byte(`{"a":"\"[[[[[["}`), false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var v any + err := Unmarshal(tt.data, &v) + gotDepthErr := errors.Is(err, errMaxDepthExceeded) + if gotDepthErr != tt.wantErr { + t.Fatalf("Unmarshal() error = %v (err=%v), want depth error = %v", gotDepthErr, err, tt.wantErr) + } + }) + } +} + +func TestUnmarshalMaxDepthExitsFast(t *testing.T) { + data := []byte(strings.Repeat("[", 500_000) + strings.Repeat("]", 500_000)) + done := make(chan error, 1) + go func() { + var v any + done <- Unmarshal(data, &v) + }() + select { + case err := <-done: + if !errors.Is(err, errMaxDepthExceeded) { + t.Fatalf("Unmarshal error = %v, want %v", err, errMaxDepthExceeded) + } + case <-time.After(500 * time.Millisecond): + t.Fatal("Unmarshal did not reject deeply nested payload promptly") + } +} + +func FuzzUnmarshalKeyComparison(f *testing.F) { + // Add seed corpus for expected valid key and the known problematic cases. + f.Add("field") + f.Add("FIELD") + f.Add("field\x00") + + f.Fuzz(func(t *testing.T, key string) { + // Generate valid JSON with the fuzzed key. + // json.Encoder ensures the key is properly escaped (e.g. \x00 becomes \u0000). + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + enc.SetEscapeHTML(false) + if err := enc.Encode(key); err != nil { + return + } + encodedKey := buf.Bytes() + jsonData := fmt.Sprintf(`{%s: "value"}`, encodedKey) + + type Target struct { + Field string `json:"field"` + } + + var got Target + if err := Unmarshal([]byte(jsonData), &got); err != nil { + return + } + + if got.Field == "value" && key != "field" { + t.Errorf("Unmarshal matched key %q (JSON: %s) to field 'field'", key, encodedKey) + } + }) +} + +func FuzzUnmarshalFieldIsolation(f *testing.F) { + // Add seed corpus. + f.Add("other", "value") + f.Add("safe", "overwrite") + f.Add("safe\x00", "attack") + + f.Fuzz(func(t *testing.T, key, val string) { + // Generate valid JSON with the fuzzed key and value. + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + enc.SetEscapeHTML(false) + if err := enc.Encode(key); err != nil { + return + } + encodedKey := buf.Bytes() + buf.Reset() + if err := enc.Encode(val); err != nil { + return + } + encodedVal := buf.Bytes() + + jsonData := fmt.Sprintf(`{"safe": "value", %s: %s}`, encodedKey, encodedVal) + + type Target struct { + Safe string `json:"safe"` + } + + var got Target + if err := Unmarshal([]byte(jsonData), &got); err != nil { + return + } + + if got.Safe != "value" && key != "safe" { + t.Errorf("Field 'safe' improperly modified by key %q (JSON: %s). Got %q, want %q", key, encodedKey, got.Safe, "value") + } + }) +} diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 5549ee1..94da6ae 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -6,58 +6,44 @@ package jsonrpc2 import ( "context" - "encoding/json" "errors" "fmt" "io" "sync" "sync/atomic" "time" + + "github.com/modelcontextprotocol/go-sdk/internal/json" + "github.com/modelcontextprotocol/go-sdk/internal/mcpgodebug" ) -// Binder builds a connection configuration. -// This may be used in servers to generate a new configuration per connection. -// ConnectionOptions itself implements Binder returning itself unmodified, to -// allow for the simple cases where no per connection information is needed. -type Binder interface { - // Bind returns the ConnectionOptions to use when establishing the passed-in - // Connection. - // - // The connection is not ready to use when Bind is called, - // but Bind may close it without reading or writing to it. - Bind(context.Context, *Connection) ConnectionOptions -} - -// A BinderFunc implements the Binder interface for a standalone Bind function. -type BinderFunc func(context.Context, *Connection) ConnectionOptions - -func (f BinderFunc) Bind(ctx context.Context, c *Connection) ConnectionOptions { - return f(ctx, c) -} - -var _ Binder = BinderFunc(nil) - -// ConnectionOptions holds the options for new connections. -type ConnectionOptions struct { - // Framer allows control over the message framing and encoding. - // If nil, HeaderFramer will be used. - Framer Framer - // Preempter allows registration of a pre-queue message handler. - // If nil, no messages will be preempted. - Preempter Preempter - // Handler is used as the queued message handler for inbound messages. - // If nil, all responses will be ErrNotHandled. - Handler Handler - // OnInternalError, if non-nil, is called with any internal errors that occur - // while serving the connection, such as protocol errors or invariant - // violations. (If nil, internal errors result in panics.) - OnInternalError func(error) -} +// nomethodnotfoundcodeinerror is a compatibility parameter that restores the +// pre-fix behavior of [processResult], where wrapped [ErrNotHandled] or +// [ErrMethodNotFound] errors returned by request handlers were not +// recognized as "method not found" signals. The original switch statement +// compared sentinel errors with ==, which never matched errors returned via +// fmt.Errorf("%w: ...", ErrNotHandled, ...) — including the ones produced +// by checkRequest. As a result the wire error response carried code 0 +// instead of code -32601. The fix uses errors.Is to recognize wrapped +// sentinels and append the method name to the message. +// +// To restore the previous behavior, set MCPGODEBUG=nomethodnotfoundcodeinerror=1. +// This option will be removed in a future SDK version. +// See the documentation for the mcpgodebug package for instructions on how +// to use it. +var nomethodnotfoundcodeinerror = mcpgodebug.Value("nomethodnotfoundcodeinerror") // Connection manages the jsonrpc2 protocol, connecting responses back to their -// calls. -// Connection is bidirectional; it does not have a designated server or client -// end. +// calls. Connection is bidirectional; it does not have a designated server or +// client end. +// +// Note that the word 'Connection' is overloaded: the mcp.Connection represents +// the bidirectional stream of messages between client an server. The +// jsonrpc2.Connection layers RPC logic on top of that stream, dispatching RPC +// handlers, and correlating requests with responses from the peer. +// +// Some of the complexity of the Connection type is grown out of its usage in +// gopls: it could probably be simplified based on our usage in MCP. type Connection struct { seq int64 // must only be accessed using atomic operations @@ -172,12 +158,12 @@ func (s *inFlightState) shuttingDown(errClosing error) error { if s.readErr != nil { // If the read side of the connection is broken, we cannot read new call // requests, and cannot read responses to our outgoing calls. - return fmt.Errorf("%w: %v", errClosing, s.readErr) + return fmt.Errorf("%w: %w", errClosing, s.readErr) } if s.writeErr != nil { // If the write side of the connection is broken, we cannot write responses // for incoming calls, and cannot write requests for outgoing calls. - return fmt.Errorf("%w: %v", errClosing, s.writeErr) + return fmt.Errorf("%w: %w", errClosing, s.writeErr) } return nil } @@ -186,12 +172,31 @@ func (s *inFlightState) shuttingDown(errClosing error) error { type incomingRequest struct { *Request // the request being processed ctx context.Context - cancel context.CancelFunc + cancel context.CancelCauseFunc } -// Bind returns the options unmodified. -func (o ConnectionOptions) Bind(context.Context, *Connection) ConnectionOptions { - return o +// Reader abstracts the transport mechanics from the JSON RPC protocol. +// A Connection reads messages from the reader it was provided on construction, +// and assumes that each call to Read fully transfers a single message, +// or returns an error. +// +// A reader is not safe for concurrent use, it is expected it will be used by +// a single Connection in a safe manner. +type Reader interface { + // Read gets the next message from the stream. + Read(context.Context) (Message, error) +} + +// Writer abstracts the transport mechanics from the JSON RPC protocol. +// A Connection writes messages using the writer it was provided on construction, +// and assumes that each call to Write fully transfers a single message, +// or returns an error. +// +// A writer must be safe for concurrent use, as writes may occur concurrently +// in practice: libraries may make calls or respond to requests asynchronously. +type Writer interface { + // Write sends a message to the stream. + Write(context.Context, Message) error } // A ConnectionConfig configures a bidirectional jsonrpc2 connection. @@ -203,12 +208,30 @@ type ConnectionConfig struct { Bind func(*Connection) Handler // required OnDone func() // optional OnInternalError func(error) // optional + + // PropagateCancellation controls whether cancellation of the context + // passed to [NewConnection] is observable by request handlers. + // + // By default (false), the connection wraps that context (see [notDone]) + // so handlers' Done channels do not fire when the connection's root + // context is cancelled. Cancellation of an in-flight handler is then + // expected to flow only through the jsonrpc2 layer's explicit channels + // (the [Preempter] reacting to the peer's cancel notification, or a + // transport read/write failure cancelling every in-flight request). + // + // Set this true when cancellation of the connection's context is itself + // a meaningful signal that handlers should react to (for example, when + // the connection is tied to a carrier that owns it and whose end means + // the request is cancelled). + PropagateCancellation bool // optional } // NewConnection creates a new [Connection] object and starts processing // incoming messages. func NewConnection(ctx context.Context, cfg ConnectionConfig) *Connection { - ctx = notDone{ctx} + if !cfg.PropagateCancellation { + ctx = notDone{ctx} + } c := &Connection{ state: inFlightState{closer: cfg.Closer}, @@ -222,59 +245,16 @@ func NewConnection(ctx context.Context, cfg ConnectionConfig) *Connection { return c } -// bindConnection creates a new connection and runs it. -// -// This is used by the Dial and Serve functions to build the actual connection. -// -// The connection is closed automatically (and its resources cleaned up) when -// the last request has completed after the underlying ReadWriteCloser breaks, -// but it may be stopped earlier by calling Close (for a clean shutdown). -func bindConnection(bindCtx context.Context, rwc io.ReadWriteCloser, binder Binder, onDone func()) *Connection { - // TODO: Should we create a new event span here? - // This will propagate cancellation from ctx; should it? - ctx := notDone{bindCtx} - - c := &Connection{ - state: inFlightState{closer: rwc}, - done: make(chan struct{}), - onDone: onDone, - } - // It's tempting to set a finalizer on c to verify that the state has gone - // idle when the connection becomes unreachable. Unfortunately, the Binder - // interface makes that unsafe: it allows the Handler to close over the - // Connection, which could create a reference cycle that would cause the - // Connection to become uncollectable. - - options := binder.Bind(bindCtx, c) - framer := options.Framer - if framer == nil { - framer = HeaderFramer() - } - c.handler = options.Handler - if c.handler == nil { - c.handler = defaultHandler{} - } - c.onInternalError = options.OnInternalError - - c.writer = framer.Writer(rwc) - reader := framer.Reader(rwc) - c.start(ctx, reader, options.Preempter) - return c -} - func (c *Connection) start(ctx context.Context, reader Reader, preempter Preempter) { c.updateInFlight(func(s *inFlightState) { select { case <-c.done: - // Bind already closed the connection; don't start a goroutine to read it. + // The connection was already closed; don't start a goroutine to read it. return default: } // The goroutine started here will continue until the underlying stream is closed. - // - // (If the Binder closed the Connection already, this should error out and - // return almost immediately.) s.reading = true go c.readIncoming(ctx, reader, preempter) }) @@ -361,19 +341,26 @@ func (c *Connection) Call(ctx context.Context, method string, params any) *Async if err := c.write(ctx, call); err != nil { // Sending failed. We will never get a response, so deliver a fake one if it // wasn't already retired by the connection breaking. - c.updateInFlight(func(s *inFlightState) { - if s.outgoingCalls[ac.id] == ac { - delete(s.outgoingCalls, ac.id) - ac.retire(&Response{ID: id, Error: err}) - } else { - // ac was already retired by the readIncoming goroutine: - // perhaps our write raced with the Read side of the connection breaking. - } - }) + c.Retire(ac, err) } return ac } +// Retire stops tracking the call, and reports err as its terminal error. +// +// Retire is safe to call multiple times: if the call is already no longer +// tracked, Retire is a no op. +func (c *Connection) Retire(ac *AsyncCall, err error) { + c.updateInFlight(func(s *inFlightState) { + if s.outgoingCalls[ac.id] == ac { + delete(s.outgoingCalls, ac.id) + ac.retire(&Response{ID: ac.id, Error: err}) + } else { + // ac was already retired elsewhere. + } + }) +} + // Async, signals that the current jsonrpc2 request may be handled // asynchronously to subsequent requests, when ctx is the request context. // @@ -424,19 +411,10 @@ type AsyncCall struct { // This can be used to cancel the call if needed. func (ac *AsyncCall) ID() ID { return ac.id } -// IsReady can be used to check if the result is already prepared. -// This is guaranteed to return true on a result for which Await has already -// returned, or a call that failed to send in the first place. -func (ac *AsyncCall) IsReady() bool { - select { - case <-ac.ready: - return true - default: - return false - } -} - // retire processes the response to the call. +// +// It is an error to call retire more than once: retire is guarded by the +// connection's outgoingCalls map. func (ac *AsyncCall) retire(response *Response) { select { case <-ac.ready: @@ -450,6 +428,9 @@ func (ac *AsyncCall) retire(response *Response) { // Await waits for (and decodes) the results of a Call. // The response will be unmarshaled from JSON into the result. +// +// If the call is cancelled due to context cancellation, the result is +// ctx.Err(). func (ac *AsyncCall) Await(ctx context.Context, result any) error { select { case <-ctx.Done(): @@ -472,12 +453,19 @@ func (ac *AsyncCall) Await(ctx context.Context, result any) error { // will not cause any messages that have not arrived yet with that ID to be // cancelled. func (c *Connection) Cancel(id ID) { + c.CancelCause(id, nil) +} + +// CancelCause is like [Connection.Cancel], but records cause as the reason +// the Context was cancelled, so that the Handle call can read it back through +// [context.Cause]. A nil cause reads as [context.Canceled]. +func (c *Connection) CancelCause(id ID, cause error) { var req *incomingRequest c.updateInFlight(func(s *inFlightState) { req = s.incomingByID[id] }) if req != nil { - req.cancel() + req.cancel(cause) } } @@ -567,6 +555,16 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter ac.retire(&Response{ID: id, Error: err}) } s.outgoingCalls = nil + + // Cancel any incoming requests still in flight: with the reader gone we + // cannot receive cancellation notifications, and likely cannot write a + // response either, so parked handlers have nothing useful left to do. + // The read error (typically io.EOF from the peer disconnecting) is the + // cause, rather than a directional ErrServerClosing/ErrClientClosing + // sentinel: at this layer the connection has no designated end. + for _, r := range s.incomingByID { + r.cancel(err) + } }) } @@ -575,7 +573,7 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter func (c *Connection) acceptRequest(ctx context.Context, msg *Request, preempter Preempter) { // In theory notifications cannot be cancelled, but we build them a cancel // context anyway. - reqCtx, cancel := context.WithCancel(ctx) + reqCtx, cancel := context.WithCancelCause(ctx) req := &incomingRequest{ Request: msg, ctx: reqCtx, @@ -676,15 +674,14 @@ func (c *Connection) handleAsync() { return } - // Only deliver to the Handler if not already canceled. + // Only deliver to the Handler if not already canceled. If the request + // was canceled with a cause (e.g. a write failure cancels every + // in-flight request), report that cause rather than the bare + // context.Canceled. if err := req.ctx.Err(); err != nil { - c.updateInFlight(func(s *inFlightState) { - if s.writeErr != nil { - // Assume that req.ctx was canceled due to s.writeErr. - // TODO(#51365): use a Context API to plumb this through req.ctx. - err = fmt.Errorf("%w: %v", ErrServerClosing, s.writeErr) - } - }) + if cause := context.Cause(req.ctx); cause != nil { + err = cause + } c.processResult("handleAsync", req, nil, err) continue } @@ -702,9 +699,7 @@ func (c *Connection) handleAsync() { // processResult processes the result of a request and, if appropriate, sends a response. func (c *Connection) processResult(from any, req *incomingRequest, result any, err error) error { - switch err { - case ErrNotHandled, ErrMethodNotFound: - // Add detail describing the unhandled method. + if nomethodnotfoundcodeinerror != "1" && (errors.Is(err, ErrNotHandled) || errors.Is(err, ErrMethodNotFound)) { err = fmt.Errorf("%w: %q", ErrMethodNotFound, req.Method) } @@ -747,7 +742,7 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e } // Cancel the request to free any associated resources. - req.cancel() + req.cancel(nil) c.updateInFlight(func(s *inFlightState) { if s.incoming == 0 { panic("jsonrpc2: processResult called when incoming count is already zero") @@ -763,22 +758,21 @@ func (c *Connection) write(ctx context.Context, msg Message) error { var err error // Fail writes immediately if the connection is shutting down. // - // TODO(rfindley): should we allow cancellation notifications through? It - // could be the case that writes can still succeed. + // Allow outgoing "notifications" forwarded by the Notify method. + // This will allow to send the cancelled notification when the client is shutting down. c.updateInFlight(func(s *inFlightState) { + if req, ok := msg.(*Request); ok && !req.IsCall() && s.outgoingNotifications > 0 { + return + } err = s.shuttingDown(ErrServerClosing) }) if err == nil { err = c.writer.Write(ctx, msg) } - // For rejected requests, we don't set the writeErr (which would break the - // connection). They can just be returned to the caller. - if errors.Is(err, ErrRejected) { - return err - } - - if err != nil && ctx.Err() == nil { + // For cancelled or rejected requests, we don't set the writeErr (which would + // break the connection). They can just be returned to the caller. + if err != nil && ctx.Err() == nil && !errors.Is(err, ErrRejected) { // The call to Write failed, and since ctx.Err() is nil we can't attribute // the failure (even indirectly) to Context cancellation. The writer appears // to be broken, and future writes are likely to also fail. @@ -791,7 +785,7 @@ func (c *Connection) write(ctx context.Context, msg Message) error { if s.writeErr == nil { s.writeErr = err for _, r := range s.incomingByID { - r.cancel() + r.cancel(fmt.Errorf("%w: %v", ErrServerClosing, err)) } } }) @@ -814,6 +808,15 @@ func (c *Connection) internalErrorf(format string, args ...any) error { } // notDone is a context.Context wrapper that returns a nil Done channel. +// +// Request handlers' contexts are derived from the connection's root context, +// which by default is wrapped in notDone so a transport-level cancellation +// does not implicitly cancel every in-flight handler. Cancellation of an +// in-flight handler is instead expected to flow only through the jsonrpc2 +// layer's explicit channels: the [Preempter] calling [Connection.Cancel] in +// response to the peer's cancel notification, or the transport itself +// failing (the read loop exits on EOF or a write fails) — both of which +// cancel every in-flight incoming request in turn. type notDone struct{ ctx context.Context } func (ic notDone) Value(key any) any { diff --git a/internal/jsonrpc2/conn_test.go b/internal/jsonrpc2/conn_test.go new file mode 100644 index 0000000..afeb4c0 --- /dev/null +++ b/internal/jsonrpc2/conn_test.go @@ -0,0 +1,38 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package jsonrpc2 + +import ( + "errors" + "io" + "testing" +) + +// TestShuttingDownWrapsWriteError verifies that when a connection shuts down +// because its write side failed, the returned error preserves the underlying +// write error in its chain so callers can classify it with errors.Is (for +// example, distinguishing io.EOF from a clean host disconnect versus a real +// failure). +func TestShuttingDownWrapsWriteError(t *testing.T) { + s := &inFlightState{writeErr: io.EOF} + err := s.shuttingDown(ErrServerClosing) + if !errors.Is(err, ErrServerClosing) { + t.Errorf("shuttingDown() error = %v, want it to wrap ErrServerClosing", err) + } + if !errors.Is(err, io.EOF) { + t.Errorf("shuttingDown() error = %v, want it to wrap io.EOF", err) + } +} + +func TestShuttingDownWrapsReadError(t *testing.T) { + s := &inFlightState{readErr: io.EOF} + err := s.shuttingDown(ErrServerClosing) + if !errors.Is(err, ErrServerClosing) { + t.Errorf("shuttingDown() error = %v, want it to wrap ErrServerClosing", err) + } + if !errors.Is(err, io.EOF) { + t.Errorf("shuttingDown() error = %v, want it to wrap io.EOF", err) + } +} diff --git a/internal/jsonrpc2/frame.go b/internal/jsonrpc2/frame.go deleted file mode 100644 index 46fcc9d..0000000 --- a/internal/jsonrpc2/frame.go +++ /dev/null @@ -1,208 +0,0 @@ -// Copyright 2018 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -package jsonrpc2 - -import ( - "bufio" - "context" - "encoding/json" - "fmt" - "io" - "strconv" - "strings" - "sync" -) - -// Reader abstracts the transport mechanics from the JSON RPC protocol. -// A Conn reads messages from the reader it was provided on construction, -// and assumes that each call to Read fully transfers a single message, -// or returns an error. -// -// A reader is not safe for concurrent use, it is expected it will be used by -// a single Conn in a safe manner. -type Reader interface { - // Read gets the next message from the stream. - Read(context.Context) (Message, error) -} - -// Writer abstracts the transport mechanics from the JSON RPC protocol. -// A Conn writes messages using the writer it was provided on construction, -// and assumes that each call to Write fully transfers a single message, -// or returns an error. -// -// A writer must be safe for concurrent use, as writes may occur concurrently -// in practice: libraries may make calls or respond to requests asynchronously. -type Writer interface { - // Write sends a message to the stream. - Write(context.Context, Message) error -} - -// Framer wraps low level byte readers and writers into jsonrpc2 message -// readers and writers. -// It is responsible for the framing and encoding of messages into wire form. -// -// TODO(rfindley): rethink the framer interface, as with JSONRPC2 batching -// there is a need for Reader and Writer to be correlated, and while the -// implementation of framing here allows that, it is not made explicit by the -// interface. -// -// Perhaps a better interface would be -// -// Frame(io.ReadWriteCloser) (Reader, Writer). -type Framer interface { - // Reader wraps a byte reader into a message reader. - Reader(io.Reader) Reader - // Writer wraps a byte writer into a message writer. - Writer(io.Writer) Writer -} - -// RawFramer returns a new Framer. -// The messages are sent with no wrapping, and rely on json decode consistency -// to determine message boundaries. -func RawFramer() Framer { return rawFramer{} } - -type rawFramer struct{} -type rawReader struct{ in *json.Decoder } -type rawWriter struct { - mu sync.Mutex - out io.Writer -} - -func (rawFramer) Reader(rw io.Reader) Reader { - return &rawReader{in: json.NewDecoder(rw)} -} - -func (rawFramer) Writer(rw io.Writer) Writer { - return &rawWriter{out: rw} -} - -func (r *rawReader) Read(ctx context.Context) (Message, error) { - select { - case <-ctx.Done(): - return nil, ctx.Err() - default: - } - var raw json.RawMessage - if err := r.in.Decode(&raw); err != nil { - return nil, err - } - msg, err := DecodeMessage(raw) - return msg, err -} - -func (w *rawWriter) Write(ctx context.Context, msg Message) error { - select { - case <-ctx.Done(): - return ctx.Err() - default: - } - - data, err := EncodeMessage(msg) - if err != nil { - return fmt.Errorf("marshaling message: %v", err) - } - - w.mu.Lock() - defer w.mu.Unlock() - _, err = w.out.Write(data) - return err -} - -// HeaderFramer returns a new Framer. -// The messages are sent with HTTP content length and MIME type headers. -// This is the format used by LSP and others. -func HeaderFramer() Framer { return headerFramer{} } - -type headerFramer struct{} -type headerReader struct{ in *bufio.Reader } -type headerWriter struct { - mu sync.Mutex - out io.Writer -} - -func (headerFramer) Reader(rw io.Reader) Reader { - return &headerReader{in: bufio.NewReader(rw)} -} - -func (headerFramer) Writer(rw io.Writer) Writer { - return &headerWriter{out: rw} -} - -func (r *headerReader) Read(ctx context.Context) (Message, error) { - select { - case <-ctx.Done(): - return nil, ctx.Err() - default: - } - - firstRead := true // to detect a clean EOF below - var contentLength int64 - // read the header, stop on the first empty line - for { - line, err := r.in.ReadString('\n') - if err != nil { - if err == io.EOF { - if firstRead && line == "" { - return nil, io.EOF // clean EOF - } - err = io.ErrUnexpectedEOF - } - return nil, fmt.Errorf("failed reading header line: %w", err) - } - firstRead = false - - line = strings.TrimSpace(line) - // check we have a header line - if line == "" { - break - } - colon := strings.IndexRune(line, ':') - if colon < 0 { - return nil, fmt.Errorf("invalid header line %q", line) - } - name, value := line[:colon], strings.TrimSpace(line[colon+1:]) - switch name { - case "Content-Length": - if contentLength, err = strconv.ParseInt(value, 10, 32); err != nil { - return nil, fmt.Errorf("failed parsing Content-Length: %v", value) - } - if contentLength <= 0 { - return nil, fmt.Errorf("invalid Content-Length: %v", contentLength) - } - default: - // ignoring unknown headers - } - } - if contentLength == 0 { - return nil, fmt.Errorf("missing Content-Length header") - } - data := make([]byte, contentLength) - _, err := io.ReadFull(r.in, data) - if err != nil { - return nil, err - } - msg, err := DecodeMessage(data) - return msg, err -} - -func (w *headerWriter) Write(ctx context.Context, msg Message) error { - select { - case <-ctx.Done(): - return ctx.Err() - default: - } - w.mu.Lock() - defer w.mu.Unlock() - - data, err := EncodeMessage(msg) - if err != nil { - return fmt.Errorf("marshaling message: %v", err) - } - _, err = fmt.Fprintf(w.out, "Content-Length: %v\r\n\r\n", len(data)) - if err == nil { - _, err = w.out.Write(data) - } - return err -} diff --git a/internal/jsonrpc2/jsonrpc2.go b/internal/jsonrpc2/jsonrpc2.go index 234e6ee..76a5116 100644 --- a/internal/jsonrpc2/jsonrpc2.go +++ b/internal/jsonrpc2/jsonrpc2.go @@ -12,17 +12,12 @@ import ( "errors" ) -var ( - // ErrIdleTimeout is returned when serving timed out waiting for new connections. - ErrIdleTimeout = errors.New("timed out waiting for new connections") - - // ErrNotHandled is returned from a Handler or Preempter to indicate it did - // not handle the request. - // - // If a Handler returns ErrNotHandled, the server replies with - // ErrMethodNotFound. - ErrNotHandled = errors.New("JSON RPC not handled") -) +// ErrNotHandled is returned from a Handler or Preempter to indicate it did +// not handle the request. +// +// If a Handler returns ErrNotHandled, the server replies with +// ErrMethodNotFound. +var ErrNotHandled = errors.New("JSON RPC not handled") // Preempter handles messages on a connection before they are queued to the main // handler. @@ -40,15 +35,6 @@ type Preempter interface { Preempt(ctx context.Context, req *Request) (result any, err error) } -// A PreempterFunc implements the Preempter interface for a standalone Preempt function. -type PreempterFunc func(ctx context.Context, req *Request) (any, error) - -func (f PreempterFunc) Preempt(ctx context.Context, req *Request) (any, error) { - return f(ctx, req) -} - -var _ Preempter = PreempterFunc(nil) - // Handler handles messages on a connection. type Handler interface { // Handle is invoked sequentially for each incoming request that has not @@ -67,16 +53,6 @@ type Handler interface { Handle(ctx context.Context, req *Request) (result any, err error) } -type defaultHandler struct{} - -func (defaultHandler) Preempt(context.Context, *Request) (any, error) { - return nil, ErrNotHandled -} - -func (defaultHandler) Handle(context.Context, *Request) (any, error) { - return nil, ErrNotHandled -} - // A HandlerFunc implements the Handler interface for a standalone Handle function. type HandlerFunc func(ctx context.Context, req *Request) (any, error) @@ -85,37 +61,3 @@ func (f HandlerFunc) Handle(ctx context.Context, req *Request) (any, error) { } var _ Handler = HandlerFunc(nil) - -// async is a small helper for operations with an asynchronous result that you -// can wait for. -type async struct { - ready chan struct{} // closed when done - firstErr chan error // 1-buffered; contains either nil or the first non-nil error -} - -func newAsync() *async { - var a async - a.ready = make(chan struct{}) - a.firstErr = make(chan error, 1) - a.firstErr <- nil - return &a -} - -func (a *async) done() { - close(a.ready) -} - -func (a *async) wait() error { - <-a.ready - err := <-a.firstErr - a.firstErr <- err - return err -} - -func (a *async) setError(err error) { - storedErr := <-a.firstErr - if storedErr == nil { - storedErr = err - } - a.firstErr <- storedErr -} diff --git a/internal/jsonrpc2/jsonrpc2_test.go b/internal/jsonrpc2/jsonrpc2_test.go deleted file mode 100644 index 8c79300..0000000 --- a/internal/jsonrpc2/jsonrpc2_test.go +++ /dev/null @@ -1,385 +0,0 @@ -// Copyright 2018 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -package jsonrpc2_test - -import ( - "context" - "encoding/json" - "fmt" - "path" - "reflect" - "testing" - "time" - - "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" -) - -var callTests = []invoker{ - call{"no_args", nil, true}, - call{"one_string", "fish", "got:fish"}, - call{"one_number", 10, "got:10"}, - call{"join", []string{"a", "b", "c"}, "a/b/c"}, - sequence{"notify", []invoker{ - notify{"set", 3}, - notify{"add", 5}, - call{"get", nil, 8}, - }}, - sequence{"preempt", []invoker{ - async{"a", "wait", "a"}, - notify{"unblock", "a"}, - collect{"a", true, false}, - }}, - sequence{"basic cancel", []invoker{ - async{"b", "wait", "b"}, - cancel{"b"}, - collect{"b", nil, true}, - }}, - sequence{"queue", []invoker{ - async{"a", "wait", "a"}, - notify{"set", 1}, - notify{"add", 2}, - notify{"add", 3}, - notify{"add", 4}, - call{"peek", nil, 0}, // accumulator will not have any adds yet - notify{"unblock", "a"}, - collect{"a", true, false}, - call{"get", nil, 10}, // accumulator now has all the adds - }}, - sequence{"fork", []invoker{ - async{"a", "fork", "a"}, - notify{"set", 1}, - notify{"add", 2}, - notify{"add", 3}, - notify{"add", 4}, - call{"get", nil, 10}, // fork will not have blocked the adds - notify{"unblock", "a"}, - collect{"a", true, false}, - }}, - sequence{"concurrent", []invoker{ - async{"a", "fork", "a"}, - notify{"unblock", "a"}, - async{"b", "fork", "b"}, - notify{"unblock", "b"}, - collect{"a", true, false}, - collect{"b", true, false}, - }}, -} - -type binder struct { - framer jsonrpc2.Framer - runTest func(*handler) -} - -type handler struct { - conn *jsonrpc2.Connection - accumulator int - waiters chan map[string]chan struct{} - calls map[string]*jsonrpc2.AsyncCall -} - -type invoker interface { - Name() string - Invoke(t *testing.T, ctx context.Context, h *handler) -} - -type notify struct { - method string - params any -} - -type call struct { - method string - params any - expect any -} - -type async struct { - name string - method string - params any -} - -type collect struct { - name string - expect any - fails bool -} - -type cancel struct { - name string -} - -type sequence struct { - name string - tests []invoker -} - -type echo call - -type cancelParams struct{ ID int64 } - -func TestConnectionRaw(t *testing.T) { - testConnection(t, jsonrpc2.RawFramer()) -} - -func TestConnectionHeader(t *testing.T) { - testConnection(t, jsonrpc2.HeaderFramer()) -} - -func testConnection(t *testing.T, framer jsonrpc2.Framer) { - ctx := context.Background() - listener, err := jsonrpc2.NetPipeListener(ctx) - if err != nil { - t.Fatal(err) - } - server := jsonrpc2.NewServer(ctx, listener, binder{framer, nil}) - defer func() { - listener.Close() - server.Wait() - }() - - for _, test := range callTests { - t.Run(test.Name(), func(t *testing.T) { - client, err := jsonrpc2.Dial(ctx, - listener.Dialer(), binder{framer, func(h *handler) { - // Sleep a little to a void a race with setting conn.writer in jsonrpc2.bindConnection. - time.Sleep(50 * time.Millisecond) - defer h.conn.Close() - test.Invoke(t, ctx, h) - if call, ok := test.(*call); ok { - // also run all simple call tests in echo mode - (*echo)(call).Invoke(t, ctx, h) - } - }}, nil) - if err != nil { - t.Fatal(err) - } - client.Wait() - }) - } -} - -func (test notify) Name() string { return test.method } -func (test notify) Invoke(t *testing.T, ctx context.Context, h *handler) { - if err := h.conn.Notify(ctx, test.method, test.params); err != nil { - t.Fatalf("%v:Notify failed: %v", test.method, err) - } -} - -func (test call) Name() string { return test.method } -func (test call) Invoke(t *testing.T, ctx context.Context, h *handler) { - results := newResults(test.expect) - if err := h.conn.Call(ctx, test.method, test.params).Await(ctx, results); err != nil { - t.Fatalf("%v:Call failed: %v", test.method, err) - } - verifyResults(t, test.method, results, test.expect) -} - -func (test echo) Invoke(t *testing.T, ctx context.Context, h *handler) { - results := newResults(test.expect) - if err := h.conn.Call(ctx, "echo", []any{test.method, test.params}).Await(ctx, results); err != nil { - t.Fatalf("%v:Echo failed: %v", test.method, err) - } - verifyResults(t, test.method, results, test.expect) -} - -func (test async) Name() string { return test.name } -func (test async) Invoke(t *testing.T, ctx context.Context, h *handler) { - h.calls[test.name] = h.conn.Call(ctx, test.method, test.params) -} - -func (test collect) Name() string { return test.name } -func (test collect) Invoke(t *testing.T, ctx context.Context, h *handler) { - o := h.calls[test.name] - results := newResults(test.expect) - err := o.Await(ctx, results) - switch { - case test.fails && err == nil: - t.Fatalf("%v:Collect was supposed to fail", test.name) - case !test.fails && err != nil: - t.Fatalf("%v:Collect failed: %v", test.name, err) - } - verifyResults(t, test.name, results, test.expect) -} - -func (test cancel) Name() string { return test.name } -func (test cancel) Invoke(t *testing.T, ctx context.Context, h *handler) { - o := h.calls[test.name] - if err := h.conn.Notify(ctx, "cancel", &cancelParams{o.ID().Raw().(int64)}); err != nil { - t.Fatalf("%v:Collect failed: %v", test.name, err) - } -} - -func (test sequence) Name() string { return test.name } -func (test sequence) Invoke(t *testing.T, ctx context.Context, h *handler) { - for _, child := range test.tests { - child.Invoke(t, ctx, h) - } -} - -// newResults makes a new empty copy of the expected type to put the results into -func newResults(expect any) any { - switch e := expect.(type) { - case []any: - var r []any - for _, v := range e { - r = append(r, reflect.New(reflect.TypeOf(v)).Interface()) - } - return r - case nil: - return nil - default: - return reflect.New(reflect.TypeOf(expect)).Interface() - } -} - -// verifyResults compares the results to the expected values -func verifyResults(t *testing.T, method string, results any, expect any) { - if expect == nil { - if results != nil { - t.Errorf("%v:Got results %+v where none expected", method, expect) - } - return - } - val := reflect.Indirect(reflect.ValueOf(results)).Interface() - if !reflect.DeepEqual(val, expect) { - t.Errorf("%v:Results are incorrect, got %+v expect %+v", method, val, expect) - } -} - -func (b binder) Bind(ctx context.Context, conn *jsonrpc2.Connection) jsonrpc2.ConnectionOptions { - h := &handler{ - conn: conn, - waiters: make(chan map[string]chan struct{}, 1), - calls: make(map[string]*jsonrpc2.AsyncCall), - } - h.waiters <- make(map[string]chan struct{}) - if b.runTest != nil { - go b.runTest(h) - } - return jsonrpc2.ConnectionOptions{ - Framer: b.framer, - Preempter: h, - Handler: h, - } -} - -func (h *handler) waiter(name string) chan struct{} { - waiters := <-h.waiters - defer func() { h.waiters <- waiters }() - waiter, found := waiters[name] - if !found { - waiter = make(chan struct{}) - waiters[name] = waiter - } - return waiter -} - -func (h *handler) Preempt(ctx context.Context, req *jsonrpc2.Request) (any, error) { - switch req.Method { - case "unblock": - var name string - if err := json.Unmarshal(req.Params, &name); err != nil { - return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err) - } - close(h.waiter(name)) - return nil, nil - case "peek": - if len(req.Params) > 0 { - return nil, fmt.Errorf("%w: expected no params", jsonrpc2.ErrInvalidParams) - } - return h.accumulator, nil - case "cancel": - var params cancelParams - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err) - } - h.conn.Cancel(jsonrpc2.Int64ID(params.ID)) - return nil, nil - default: - return nil, jsonrpc2.ErrNotHandled - } -} - -func (h *handler) Handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { - switch req.Method { - case "no_args": - if len(req.Params) > 0 { - return nil, fmt.Errorf("%w: expected no params", jsonrpc2.ErrInvalidParams) - } - return true, nil - case "one_string": - var v string - if err := json.Unmarshal(req.Params, &v); err != nil { - return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err) - } - return "got:" + v, nil - case "one_number": - var v int - if err := json.Unmarshal(req.Params, &v); err != nil { - return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err) - } - return fmt.Sprintf("got:%d", v), nil - case "set": - var v int - if err := json.Unmarshal(req.Params, &v); err != nil { - return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err) - } - h.accumulator = v - return nil, nil - case "add": - var v int - if err := json.Unmarshal(req.Params, &v); err != nil { - return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err) - } - h.accumulator += v - return nil, nil - case "get": - if len(req.Params) > 0 { - return nil, fmt.Errorf("%w: expected no params", jsonrpc2.ErrInvalidParams) - } - return h.accumulator, nil - case "join": - var v []string - if err := json.Unmarshal(req.Params, &v); err != nil { - return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err) - } - return path.Join(v...), nil - case "echo": - var v []any - if err := json.Unmarshal(req.Params, &v); err != nil { - return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err) - } - var result any - err := h.conn.Call(ctx, v[0].(string), v[1]).Await(ctx, &result) - return result, err - case "wait": - var name string - if err := json.Unmarshal(req.Params, &name); err != nil { - return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err) - } - select { - case <-h.waiter(name): - return true, nil - case <-ctx.Done(): - return nil, ctx.Err() - } - case "fork": - var name string - if err := json.Unmarshal(req.Params, &name); err != nil { - return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err) - } - jsonrpc2.Async(ctx) - waitFor := h.waiter(name) - select { - case <-waitFor: - return true, nil - case <-ctx.Done(): - return nil, ctx.Err() - } - default: - return nil, jsonrpc2.ErrNotHandled - } -} diff --git a/internal/jsonrpc2/messages.go b/internal/jsonrpc2/messages.go index 791e698..8b96770 100644 --- a/internal/jsonrpc2/messages.go +++ b/internal/jsonrpc2/messages.go @@ -5,9 +5,12 @@ package jsonrpc2 import ( + "bytes" "encoding/json" "errors" "fmt" + + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" ) // ID is a Request identifier, which is defined by the spec to be a string, integer, or null. @@ -145,9 +148,9 @@ func toWireError(err error) *WireError { func EncodeMessage(msg Message) ([]byte, error) { wire := wireCombined{VersionTag: wireVersion} msg.marshal(&wire) - data, err := json.Marshal(&wire) + data, err := jsonMarshal(&wire) if err != nil { - return data, fmt.Errorf("marshaling jsonrpc message: %w", err) + return nil, fmt.Errorf("marshaling jsonrpc message: %w", err) } return data, nil } @@ -158,16 +161,31 @@ func EncodeMessage(msg Message) ([]byte, error) { func EncodeIndent(msg Message, prefix, indent string) ([]byte, error) { wire := wireCombined{VersionTag: wireVersion} msg.marshal(&wire) - data, err := json.MarshalIndent(&wire, prefix, indent) - if err != nil { - return data, fmt.Errorf("marshaling jsonrpc message: %w", err) + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + enc.SetEscapeHTML(false) + enc.SetIndent(prefix, indent) + if err := enc.Encode(&wire); err != nil { + return nil, fmt.Errorf("marshaling jsonrpc message: %w", err) } - return data, nil + return bytes.TrimRight(buf.Bytes(), "\n"), nil +} + +// wireDecode is the decode form of [wireCombined]. Method is a [json.RawMessage] +// so we can tell whether the "method" key was present on the wire, including +// when its value is the empty string (see go-sdk#976). +type wireDecode struct { + VersionTag string `json:"jsonrpc"` + ID any `json:"id,omitempty"` + Method json.RawMessage `json:"method"` + Params json.RawMessage `json:"params,omitempty"` + Result json.RawMessage `json:"result,omitempty"` + Error *WireError `json:"error,omitempty"` } func DecodeMessage(data []byte) (Message, error) { - msg := wireCombined{} - if err := json.Unmarshal(data, &msg); err != nil { + msg := wireDecode{} + if err := internaljson.Unmarshal(data, &msg); err != nil { return nil, fmt.Errorf("unmarshaling jsonrpc message: %w", err) } if msg.VersionTag != wireVersion { @@ -177,15 +195,19 @@ func DecodeMessage(data []byte) (Message, error) { if err != nil { return nil, err } - if msg.Method != "" { - // has a method, must be a call + if len(msg.Method) > 0 { + // The "method" key was present. Decode its value (including ""). + var method string + if err := internaljson.Unmarshal(msg.Method, &method); err != nil { + return nil, fmt.Errorf("unmarshaling jsonrpc message: %w", err) + } return &Request{ - Method: msg.Method, + Method: method, ID: id, Params: msg.Params, }, nil } - // no method, should be a response + // no method key, should be a response if !id.IsValid() { return nil, ErrInvalidRequest } @@ -204,9 +226,21 @@ func marshalToRaw(obj any) (json.RawMessage, error) { if obj == nil { return nil, nil } - data, err := json.Marshal(obj) + data, err := jsonMarshal(obj) if err != nil { return nil, err } return json.RawMessage(data), nil } + +// jsonMarshal marshals obj to JSON like json.Marshal but without HTML escaping. +func jsonMarshal(obj any) ([]byte, error) { + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + enc.SetEscapeHTML(false) + if err := enc.Encode(obj); err != nil { + return nil, err + } + // json.Encoder.Encode adds a trailing newline. Trim it to be consistent with json.Marshal. + return bytes.TrimRight(buf.Bytes(), "\n"), nil +} diff --git a/internal/jsonrpc2/net.go b/internal/jsonrpc2/net.go deleted file mode 100644 index 05db062..0000000 --- a/internal/jsonrpc2/net.go +++ /dev/null @@ -1,138 +0,0 @@ -// Copyright 2018 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -package jsonrpc2 - -import ( - "context" - "io" - "net" - "os" -) - -// This file contains implementations of the transport primitives that use the standard network -// package. - -// NetListenOptions is the optional arguments to the NetListen function. -type NetListenOptions struct { - NetListenConfig net.ListenConfig - NetDialer net.Dialer -} - -// NetListener returns a new Listener that listens on a socket using the net package. -func NetListener(ctx context.Context, network, address string, options NetListenOptions) (Listener, error) { - ln, err := options.NetListenConfig.Listen(ctx, network, address) - if err != nil { - return nil, err - } - return &netListener{net: ln}, nil -} - -// netListener is the implementation of Listener for connections made using the net package. -type netListener struct { - net net.Listener -} - -// Accept blocks waiting for an incoming connection to the listener. -func (l *netListener) Accept(context.Context) (io.ReadWriteCloser, error) { - return l.net.Accept() -} - -// Close will cause the listener to stop listening. It will not close any connections that have -// already been accepted. -func (l *netListener) Close() error { - addr := l.net.Addr() - err := l.net.Close() - if addr.Network() == "unix" { - rerr := os.Remove(addr.String()) - if rerr != nil && err == nil { - err = rerr - } - } - return err -} - -// Dialer returns a dialer that can be used to connect to the listener. -func (l *netListener) Dialer() Dialer { - return NetDialer(l.net.Addr().Network(), l.net.Addr().String(), net.Dialer{}) -} - -// NetDialer returns a Dialer using the supplied standard network dialer. -func NetDialer(network, address string, nd net.Dialer) Dialer { - return &netDialer{ - network: network, - address: address, - dialer: nd, - } -} - -type netDialer struct { - network string - address string - dialer net.Dialer -} - -func (n *netDialer) Dial(ctx context.Context) (io.ReadWriteCloser, error) { - return n.dialer.DialContext(ctx, n.network, n.address) -} - -// NetPipeListener returns a new Listener that listens using net.Pipe. -// It is only possibly to connect to it using the Dialer returned by the -// Dialer method, each call to that method will generate a new pipe the other -// side of which will be returned from the Accept call. -func NetPipeListener(ctx context.Context) (Listener, error) { - return &netPiper{ - done: make(chan struct{}), - dialed: make(chan io.ReadWriteCloser), - }, nil -} - -// netPiper is the implementation of Listener build on top of net.Pipes. -type netPiper struct { - done chan struct{} - dialed chan io.ReadWriteCloser -} - -// Accept blocks waiting for an incoming connection to the listener. -func (l *netPiper) Accept(context.Context) (io.ReadWriteCloser, error) { - // Block until the pipe is dialed or the listener is closed, - // preferring the latter if already closed at the start of Accept. - select { - case <-l.done: - return nil, net.ErrClosed - default: - } - select { - case rwc := <-l.dialed: - return rwc, nil - case <-l.done: - return nil, net.ErrClosed - } -} - -// Close will cause the listener to stop listening. It will not close any connections that have -// already been accepted. -func (l *netPiper) Close() error { - // unblock any accept calls that are pending - close(l.done) - return nil -} - -func (l *netPiper) Dialer() Dialer { - return l -} - -func (l *netPiper) Dial(ctx context.Context) (io.ReadWriteCloser, error) { - client, server := net.Pipe() - - select { - case l.dialed <- server: - return client, nil - - case <-l.done: - client.Close() - server.Close() - return nil, net.ErrClosed - } -} diff --git a/internal/jsonrpc2/serve.go b/internal/jsonrpc2/serve.go deleted file mode 100644 index 424163a..0000000 --- a/internal/jsonrpc2/serve.go +++ /dev/null @@ -1,330 +0,0 @@ -// Copyright 2020 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -package jsonrpc2 - -import ( - "context" - "fmt" - "io" - "runtime" - "sync" - "sync/atomic" - "time" -) - -// Listener is implemented by protocols to accept new inbound connections. -type Listener interface { - // Accept accepts an inbound connection to a server. - // It blocks until either an inbound connection is made, or the listener is closed. - Accept(context.Context) (io.ReadWriteCloser, error) - - // Close closes the listener. - // Any blocked Accept or Dial operations will unblock and return errors. - Close() error - - // Dialer returns a dialer that can be used to connect to this listener - // locally. - // If a listener does not implement this it will return nil. - Dialer() Dialer -} - -// Dialer is used by clients to dial a server. -type Dialer interface { - // Dial returns a new communication byte stream to a listening server. - Dial(ctx context.Context) (io.ReadWriteCloser, error) -} - -// Server is a running server that is accepting incoming connections. -type Server struct { - listener Listener - binder Binder - async *async - - shutdownOnce sync.Once - closing int32 // atomic: set to nonzero when Shutdown is called -} - -// Dial uses the dialer to make a new connection, wraps the returned -// reader and writer using the framer to make a stream, and then builds -// a connection on top of that stream using the binder. -// -// The returned Connection will operate independently using the Preempter and/or -// Handler provided by the Binder, and will release its own resources when the -// connection is broken, but the caller may Close it earlier to stop accepting -// (or sending) new requests. -// -// If non-nil, the onDone function is called when the connection is closed. -func Dial(ctx context.Context, dialer Dialer, binder Binder, onDone func()) (*Connection, error) { - // dial a server - rwc, err := dialer.Dial(ctx) - if err != nil { - return nil, err - } - return bindConnection(ctx, rwc, binder, onDone), nil -} - -// NewServer starts a new server listening for incoming connections and returns -// it. -// This returns a fully running and connected server, it does not block on -// the listener. -// You can call Wait to block on the server, or Shutdown to get the sever to -// terminate gracefully. -// To notice incoming connections, use an intercepting Binder. -func NewServer(ctx context.Context, listener Listener, binder Binder) *Server { - server := &Server{ - listener: listener, - binder: binder, - async: newAsync(), - } - go server.run(ctx) - return server -} - -// Wait returns only when the server has shut down. -func (s *Server) Wait() error { - return s.async.wait() -} - -// Shutdown informs the server to stop accepting new connections. -func (s *Server) Shutdown() { - s.shutdownOnce.Do(func() { - atomic.StoreInt32(&s.closing, 1) - s.listener.Close() - }) -} - -// run accepts incoming connections from the listener, -// If IdleTimeout is non-zero, run exits after there are no clients for this -// duration, otherwise it exits only on error. -func (s *Server) run(ctx context.Context) { - defer s.async.done() - - var activeConns sync.WaitGroup - for { - rwc, err := s.listener.Accept(ctx) - if err != nil { - // Only Shutdown closes the listener. If we get an error after Shutdown is - // called, assume that was the cause and don't report the error; - // otherwise, report the error in case it is unexpected. - if atomic.LoadInt32(&s.closing) == 0 { - s.async.setError(err) - } - // We are done generating new connections for good. - break - } - - // A new inbound connection. - activeConns.Add(1) - _ = bindConnection(ctx, rwc, s.binder, activeConns.Done) // unregisters itself when done - } - activeConns.Wait() -} - -// NewIdleListener wraps a listener with an idle timeout. -// -// When there are no active connections for at least the timeout duration, -// calls to Accept will fail with ErrIdleTimeout. -// -// A connection is considered inactive as soon as its Close method is called. -func NewIdleListener(timeout time.Duration, wrap Listener) Listener { - l := &idleListener{ - wrapped: wrap, - timeout: timeout, - active: make(chan int, 1), - timedOut: make(chan struct{}), - idleTimer: make(chan *time.Timer, 1), - } - l.idleTimer <- time.AfterFunc(l.timeout, l.timerExpired) - return l -} - -type idleListener struct { - wrapped Listener - timeout time.Duration - - // Only one of these channels is receivable at any given time. - active chan int // count of active connections; closed when Close is called if not timed out - timedOut chan struct{} // closed when the idle timer expires - idleTimer chan *time.Timer // holds the timer only when idle -} - -// Accept accepts an incoming connection. -// -// If an incoming connection is accepted concurrent to the listener being closed -// due to idleness, the new connection is immediately closed. -func (l *idleListener) Accept(ctx context.Context) (io.ReadWriteCloser, error) { - rwc, err := l.wrapped.Accept(ctx) - - select { - case n, ok := <-l.active: - if err != nil { - if ok { - l.active <- n - } - return nil, err - } - if ok { - l.active <- n + 1 - } else { - // l.wrapped.Close Close has been called, but Accept returned a - // connection. This race can occur with concurrent Accept and Close calls - // with any net.Listener, and it is benign: since the listener was closed - // explicitly, it can't have also timed out. - } - return l.newConn(rwc), nil - - case <-l.timedOut: - if err == nil { - // Keeping the connection open would leave the listener simultaneously - // active and closed due to idleness, which would be contradictory and - // confusing. Close the connection and pretend that it never happened. - rwc.Close() - } else { - // In theory the timeout could have raced with an unrelated error return - // from Accept. However, ErrIdleTimeout is arguably still valid (since we - // would have closed due to the timeout independent of the error), and the - // harm from returning a spurious ErrIdleTimeout is negligible anyway. - } - return nil, ErrIdleTimeout - - case timer := <-l.idleTimer: - if err != nil { - // The idle timer doesn't run until it receives itself from the idleTimer - // channel, so it can't have called l.wrapped.Close yet and thus err can't - // be ErrIdleTimeout. Leave the idle timer as it was and return whatever - // error we got. - l.idleTimer <- timer - return nil, err - } - - if !timer.Stop() { - // Failed to stop the timer — the timer goroutine is in the process of - // firing. Send the timer back to the timer goroutine so that it can - // safely close the timedOut channel, and then wait for the listener to - // actually be closed before we return ErrIdleTimeout. - l.idleTimer <- timer - rwc.Close() - <-l.timedOut - return nil, ErrIdleTimeout - } - - l.active <- 1 - return l.newConn(rwc), nil - } -} - -func (l *idleListener) Close() error { - select { - case _, ok := <-l.active: - if ok { - close(l.active) - } - - case <-l.timedOut: - // Already closed by the timer; take care not to double-close if the caller - // only explicitly invokes this Close method once, since the io.Closer - // interface explicitly leaves doubled Close calls undefined. - return ErrIdleTimeout - - case timer := <-l.idleTimer: - if !timer.Stop() { - // Couldn't stop the timer. It shouldn't take long to run, so just wait - // (so that the Listener is guaranteed to be closed before we return) - // and pretend that this call happened afterward. - // That way we won't leak any timers or goroutines when Close returns. - l.idleTimer <- timer - <-l.timedOut - return ErrIdleTimeout - } - close(l.active) - } - - return l.wrapped.Close() -} - -func (l *idleListener) Dialer() Dialer { - return l.wrapped.Dialer() -} - -func (l *idleListener) timerExpired() { - select { - case n, ok := <-l.active: - if ok { - panic(fmt.Sprintf("jsonrpc2: idleListener idle timer fired with %d connections still active", n)) - } else { - panic("jsonrpc2: Close finished with idle timer still running") - } - - case <-l.timedOut: - panic("jsonrpc2: idleListener idle timer fired more than once") - - case <-l.idleTimer: - // The timer for this very call! - } - - // Close the Listener with all channels still blocked to ensure that this call - // to l.wrapped.Close doesn't race with the one in l.Close. - defer close(l.timedOut) - l.wrapped.Close() -} - -func (l *idleListener) connClosed() { - select { - case n, ok := <-l.active: - if !ok { - // l is already closed, so it can't close due to idleness, - // and we don't need to track the number of active connections any more. - return - } - n-- - if n == 0 { - l.idleTimer <- time.AfterFunc(l.timeout, l.timerExpired) - } else { - l.active <- n - } - - case <-l.timedOut: - panic("jsonrpc2: idleListener idle timer fired before last active connection was closed") - - case <-l.idleTimer: - panic("jsonrpc2: idleListener idle timer active before last active connection was closed") - } -} - -type idleListenerConn struct { - wrapped io.ReadWriteCloser - l *idleListener - closeOnce sync.Once -} - -func (l *idleListener) newConn(rwc io.ReadWriteCloser) *idleListenerConn { - c := &idleListenerConn{ - wrapped: rwc, - l: l, - } - - // A caller that forgets to call Close may disrupt the idleListener's - // accounting, even though the file descriptor for the underlying connection - // may eventually be garbage-collected anyway. - // - // Set a (best-effort) finalizer to verify that a Close call always occurs. - // (We will clear the finalizer explicitly in Close.) - runtime.SetFinalizer(c, func(c *idleListenerConn) { - panic("jsonrpc2: IdleListener connection became unreachable without a call to Close") - }) - - return c -} - -func (c *idleListenerConn) Read(p []byte) (int, error) { return c.wrapped.Read(p) } -func (c *idleListenerConn) Write(p []byte) (int, error) { return c.wrapped.Write(p) } - -func (c *idleListenerConn) Close() error { - defer c.closeOnce.Do(func() { - c.l.connClosed() - runtime.SetFinalizer(c, nil) - }) - return c.wrapped.Close() -} diff --git a/internal/jsonrpc2/serve_test.go b/internal/jsonrpc2/serve_test.go deleted file mode 100644 index 282a684..0000000 --- a/internal/jsonrpc2/serve_test.go +++ /dev/null @@ -1,354 +0,0 @@ -// Copyright 2020 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -package jsonrpc2_test - -import ( - "context" - "errors" - "fmt" - "runtime" - "runtime/debug" - "testing" - "time" - - "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" -) - -// needsLocalhostNet skips t if networking does not work for ports opened -// with "localhost". -// forked from golang.org/x/tools/internal/testenv. -func needsLocalhostNet(t testing.TB) { - switch runtime.GOOS { - case "js", "wasip1": - t.Skipf(`Listening on "localhost" fails on %s; see https://go.dev/issue/59718`, runtime.GOOS) - } -} - -func TestIdleTimeout(t *testing.T) { - needsLocalhostNet(t) - - // Use a panicking time.AfterFunc instead of context.WithTimeout so that we - // get a goroutine dump on failure. We expect the test to take on the order of - // a few tens of milliseconds at most, so 10s should be several orders of - // magnitude of headroom. - timer := time.AfterFunc(10*time.Second, func() { - debug.SetTraceback("all") - panic("TestIdleTimeout deadlocked") - }) - defer timer.Stop() - - ctx := context.Background() - - try := func(d time.Duration) (longEnough bool) { - listener, err := jsonrpc2.NetListener(ctx, "tcp", "localhost:0", jsonrpc2.NetListenOptions{}) - if err != nil { - t.Fatal(err) - } - - idleStart := time.Now() - listener = jsonrpc2.NewIdleListener(d, listener) - defer listener.Close() - - server := jsonrpc2.NewServer(ctx, listener, jsonrpc2.ConnectionOptions{}) - - // Exercise some connection/disconnection patterns, and then assert that when - // our timer fires, the server exits. - conn1, err := jsonrpc2.Dial(ctx, listener.Dialer(), jsonrpc2.ConnectionOptions{}, nil) - if err != nil { - if since := time.Since(idleStart); since < d { - t.Fatalf("conn1 failed to connect after %v: %v", since, err) - } - t.Log("jsonrpc2.Dial:", err) - return false // Took to long to dial, so the failure could have been due to the idle timeout. - } - // On the server side, Accept can race with the connection timing out. - // Send a call and wait for the response to ensure that the connection was - // actually fully accepted. - ac := conn1.Call(ctx, "ping", nil) - if err := ac.Await(ctx, nil); !errors.Is(err, jsonrpc2.ErrMethodNotFound) { - if since := time.Since(idleStart); since < d { - t.Fatalf("conn1 broken after %v: %v", since, err) - } - t.Log(`conn1.Call(ctx, "ping", nil):`, err) - conn1.Close() - return false - } - - // Since conn1 was successfully accepted and remains open, the server is - // definitely non-idle. Dialing another simultaneous connection should - // succeed. - conn2, err := jsonrpc2.Dial(ctx, listener.Dialer(), jsonrpc2.ConnectionOptions{}, nil) - if err != nil { - conn1.Close() - t.Fatalf("conn2 failed to connect while non-idle after %v: %v", time.Since(idleStart), err) - return false - } - // Ensure that conn2 is also accepted on the server side before we close - // conn1. Otherwise, the connection can appear idle if the server processes - // the closure of conn1 and the idle timeout before it finally notices conn2 - // in the accept queue. - // (That failure mode may explain the failure noted in - // https://go.dev/issue/49387#issuecomment-1303979877.) - ac = conn2.Call(ctx, "ping", nil) - if err := ac.Await(ctx, nil); !errors.Is(err, jsonrpc2.ErrMethodNotFound) { - t.Fatalf("conn2 broken while non-idle after %v: %v", time.Since(idleStart), err) - } - - if err := conn1.Close(); err != nil { - t.Fatalf("conn1.Close failed with error: %v", err) - } - idleStart = time.Now() - if err := conn2.Close(); err != nil { - t.Fatalf("conn2.Close failed with error: %v", err) - } - - conn3, err := jsonrpc2.Dial(ctx, listener.Dialer(), jsonrpc2.ConnectionOptions{}, nil) - if err != nil { - if since := time.Since(idleStart); since < d { - t.Fatalf("conn3 failed to connect after %v: %v", since, err) - } - t.Log("jsonrpc2.Dial:", err) - return false // Took to long to dial, so the failure could have been due to the idle timeout. - } - - ac = conn3.Call(ctx, "ping", nil) - if err := ac.Await(ctx, nil); !errors.Is(err, jsonrpc2.ErrMethodNotFound) { - if since := time.Since(idleStart); since < d { - t.Fatalf("conn3 broken after %v: %v", since, err) - } - t.Log(`conn3.Call(ctx, "ping", nil):`, err) - conn3.Close() - return false - } - - idleStart = time.Now() - if err := conn3.Close(); err != nil { - t.Fatalf("conn3.Close failed with error: %v", err) - } - - serverError := server.Wait() - - if !errors.Is(serverError, jsonrpc2.ErrIdleTimeout) { - t.Errorf("run() returned error %v, want %v", serverError, jsonrpc2.ErrIdleTimeout) - } - if since := time.Since(idleStart); since < d { - t.Errorf("server shut down after %v idle; want at least %v", since, d) - } - return true - } - - d := 1 * time.Millisecond - for { - t.Logf("testing with idle timeout %v", d) - if !try(d) { - d *= 2 - continue - } - break - } -} - -type msg struct { - Msg string -} - -type fakeHandler struct{} - -func (fakeHandler) Handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { - switch req.Method { - case "ping": - return &msg{"pong"}, nil - default: - return nil, jsonrpc2.ErrNotHandled - } -} - -func TestServe(t *testing.T) { - ctx := context.Background() - - tests := []struct { - name string - factory func(context.Context, testing.TB) (jsonrpc2.Listener, error) - }{ - {"tcp", func(ctx context.Context, t testing.TB) (jsonrpc2.Listener, error) { - needsLocalhostNet(t) - return jsonrpc2.NetListener(ctx, "tcp", "localhost:0", jsonrpc2.NetListenOptions{}) - }}, - {"pipe", func(ctx context.Context, t testing.TB) (jsonrpc2.Listener, error) { - return jsonrpc2.NetPipeListener(ctx) - }}, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - fake, err := test.factory(ctx, t) - if err != nil { - t.Fatal(err) - } - conn, shutdown, err := newFake(t, ctx, fake) - if err != nil { - t.Fatal(err) - } - defer shutdown() - var got msg - if err := conn.Call(ctx, "ping", &msg{"ting"}).Await(ctx, &got); err != nil { - t.Fatal(err) - } - if want := "pong"; got.Msg != want { - t.Errorf("conn.Call(...): returned %q, want %q", got, want) - } - }) - } -} - -func newFake(t *testing.T, ctx context.Context, l jsonrpc2.Listener) (*jsonrpc2.Connection, func(), error) { - server := jsonrpc2.NewServer(ctx, l, jsonrpc2.ConnectionOptions{ - Handler: fakeHandler{}, - }) - - client, err := jsonrpc2.Dial(ctx, - l.Dialer(), - jsonrpc2.ConnectionOptions{ - Handler: fakeHandler{}, - }, nil) - if err != nil { - return nil, nil, err - } - return client, func() { - if err := l.Close(); err != nil { - t.Fatal(err) - } - if err := client.Close(); err != nil { - t.Fatal(err) - } - server.Wait() - }, nil -} - -// TestIdleListenerAcceptCloseRace checks for the Accept/Close race fixed in CL 388597. -// -// (A bug in the idleListener implementation caused a successful Accept to block -// on sending to a background goroutine that could have already exited.) -func TestIdleListenerAcceptCloseRace(t *testing.T) { - ctx := context.Background() - - n := 10 - - // Each iteration of the loop appears to take around a millisecond, so to - // avoid spurious failures we'll set the watchdog for three orders of - // magnitude longer. When the bug was present, this reproduced the deadlock - // reliably on a Linux workstation when run with -count=100, which should be - // frequent enough to show up on the Go build dashboard if it regresses. - watchdog := time.Duration(n) * 1000 * time.Millisecond - timer := time.AfterFunc(watchdog, func() { - debug.SetTraceback("all") - panic(fmt.Sprintf("%s deadlocked after %v", t.Name(), watchdog)) - }) - defer timer.Stop() - - for ; n > 0; n-- { - listener, err := jsonrpc2.NetPipeListener(ctx) - if err != nil { - t.Fatal(err) - } - listener = jsonrpc2.NewIdleListener(24*time.Hour, listener) - - done := make(chan struct{}) - go func() { - conn, err := jsonrpc2.Dial(ctx, listener.Dialer(), jsonrpc2.ConnectionOptions{}, nil) - listener.Close() - if err == nil { - conn.Close() - } - close(done) - }() - - // Accept may return a non-nil error if Close closes the underlying network - // connection before the wrapped Accept call unblocks. However, it must not - // deadlock! - c, err := listener.Accept(ctx) - if err == nil { - c.Close() - } - <-done - } -} - -// TestCloseCallRace checks for a race resulting in a deadlock when a Call on -// one side of the connection races with a Close (or otherwise broken -// connection) initiated from the other side. -// -// (The Call method was waiting for a result from the Read goroutine to -// determine which error value to return, but the Read goroutine was waiting for -// in-flight calls to complete before reporting that result.) -func TestCloseCallRace(t *testing.T) { - ctx := context.Background() - n := 10 - - watchdog := time.Duration(n) * 1000 * time.Millisecond - timer := time.AfterFunc(watchdog, func() { - debug.SetTraceback("all") - panic(fmt.Sprintf("%s deadlocked after %v", t.Name(), watchdog)) - }) - defer timer.Stop() - - for ; n > 0; n-- { - listener, err := jsonrpc2.NetPipeListener(ctx) - if err != nil { - t.Fatal(err) - } - - pokec := make(chan *jsonrpc2.AsyncCall, 1) - - s := jsonrpc2.NewServer(ctx, listener, jsonrpc2.BinderFunc(func(_ context.Context, srvConn *jsonrpc2.Connection) jsonrpc2.ConnectionOptions { - h := jsonrpc2.HandlerFunc(func(ctx context.Context, _ *jsonrpc2.Request) (any, error) { - // Start a concurrent call from the server to the client. - // The point of this test is to ensure this doesn't deadlock - // if the client shuts down the connection concurrently. - // - // The racing Call may or may not receive a response: it should get a - // response if it is sent before the client closes the connection, and - // it should fail with some kind of "connection closed" error otherwise. - go func() { - pokec <- srvConn.Call(ctx, "poke", nil) - }() - - return &msg{"pong"}, nil - }) - return jsonrpc2.ConnectionOptions{Handler: h} - })) - - dialConn, err := jsonrpc2.Dial(ctx, listener.Dialer(), jsonrpc2.ConnectionOptions{}, nil) - if err != nil { - listener.Close() - s.Wait() - t.Fatal(err) - } - - // Calling any method on the server should provoke it to asynchronously call - // us back. While it is starting that call, we will close the connection. - if err := dialConn.Call(ctx, "ping", nil).Await(ctx, nil); err != nil { - t.Error(err) - } - if err := dialConn.Close(); err != nil { - t.Error(err) - } - - // Ensure that the Call on the server side did not block forever when the - // connection closed. - pokeCall := <-pokec - if err := pokeCall.Await(ctx, nil); err == nil { - t.Errorf("unexpected nil error from server-initited call") - } else if errors.Is(err, jsonrpc2.ErrMethodNotFound) { - // The call completed before the Close reached the handler. - } else { - // The error was something else. - t.Logf("server-initiated call completed with expected error: %v", err) - } - - listener.Close() - s.Wait() - } -} diff --git a/internal/jsonrpc2/wire.go b/internal/jsonrpc2/wire.go index 8be2872..499b026 100644 --- a/internal/jsonrpc2/wire.go +++ b/internal/jsonrpc2/wire.go @@ -28,9 +28,6 @@ var ( // The following errors are not part of the json specification, but // compliant extensions specific to this implementation. - // ErrServerOverloaded is returned when a message was refused due to a - // server being temporarily unable to accept any new messages. - ErrServerOverloaded = NewError(-32000, "overloaded") // ErrUnknown should be used for all non coded errors. ErrUnknown = NewError(-32001, "unknown error") // ErrServerClosing is returned for calls that arrive while the server is closing. @@ -47,7 +44,7 @@ var ( // Such failures do not indicate that the connection is broken, but rather // should be returned to the caller to indicate that the specific request is // invalid in the current context. - ErrRejected = NewError(-32004, "rejected by transport") + ErrRejected = NewError(-32005, "rejected by transport") ) const wireVersion = "2.0" @@ -90,7 +87,7 @@ func (err *WireError) Error() string { func (err *WireError) Is(other error) bool { w, ok := other.(*WireError) - if !ok { + if !ok || w == nil { return false } return err.Code == w.Code diff --git a/internal/jsonrpc2/wire_test.go b/internal/jsonrpc2/wire_test.go index 6f576f0..ef47144 100644 --- a/internal/jsonrpc2/wire_test.go +++ b/internal/jsonrpc2/wire_test.go @@ -7,6 +7,7 @@ package jsonrpc2_test import ( "bytes" "encoding/json" + "errors" "reflect" "testing" @@ -102,6 +103,74 @@ func newResponse(id any, result any, rerr error) jsonrpc2.Message { return msg } +// go-sdk#976: empty method must decode as a request (encode omits empty method). +func TestDecodeEmptyMethodRequest(t *testing.T) { + encoded := []byte(`{"jsonrpc":"2.0","id":5,"method":"","params":{}}`) + msg, err := jsonrpc2.DecodeMessage(encoded) + if err != nil { + t.Fatal(err) + } + req, ok := msg.(*jsonrpc2.Request) + if !ok { + t.Fatalf("message type = %T, want *jsonrpc2.Request", msg) + } + if req.Method != "" { + t.Errorf("Method = %q, want empty string", req.Method) + } + if req.ID != jsonrpc2.Int64ID(5) { + t.Errorf("ID = %v, want 5", req.ID.Raw()) + } + if !req.IsCall() { + t.Error("empty method with id=5 should be a call") + } +} + +func TestDecodeResponseUnchanged(t *testing.T) { + encoded := []byte(`{"jsonrpc":"2.0","id":2,"result":{}}`) + msg, err := jsonrpc2.DecodeMessage(encoded) + if err != nil { + t.Fatal(err) + } + if _, ok := msg.(*jsonrpc2.Response); !ok { + t.Fatalf("message type = %T, want *jsonrpc2.Response", msg) + } +} + +// Messages with an id but no "method" key are responses, not malformed requests. +func TestDecodeIDOnlyMessageIsResponse(t *testing.T) { + encoded := []byte(`{"jsonrpc":"2.0","id":5}`) + msg, err := jsonrpc2.DecodeMessage(encoded) + if err != nil { + t.Fatal(err) + } + if _, ok := msg.(*jsonrpc2.Response); !ok { + t.Fatalf("message type = %T, want *jsonrpc2.Response", msg) + } +} + +// TestWireErrorIs checks the errors.Is behavior of WireError, including that +// comparing against a typed-nil *WireError does not panic. +func TestWireErrorIs(t *testing.T) { + err := jsonrpc2.NewError(-32600, "invalid request") + for _, test := range []struct { + name string + target error + want bool + }{ + {name: "same code", target: jsonrpc2.NewError(-32600, "other message"), want: true}, + {name: "different code", target: jsonrpc2.NewError(-32601, "invalid request"), want: false}, + {name: "typed nil WireError", target: (*jsonrpc2.WireError)(nil), want: false}, + {name: "untyped nil", target: nil, want: false}, + {name: "other error type", target: errors.New("invalid request"), want: false}, + } { + t.Run(test.name, func(t *testing.T) { + if got := errors.Is(err, test.target); got != test.want { + t.Errorf("errors.Is(%v, %v) = %v, want %v", err, test.target, got, test.want) + } + }) + } +} + func checkJSON(t *testing.T, got, want []byte) { // compare the compact form, to allow for formatting differences g := &bytes.Buffer{} diff --git a/internal/mcpgodebug/mcpgodebug.go b/internal/mcpgodebug/mcpgodebug.go new file mode 100644 index 0000000..7f8f7ca --- /dev/null +++ b/internal/mcpgodebug/mcpgodebug.go @@ -0,0 +1,52 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +// Package mcpgodebug provides a mechanism to configure compatibility parameters +// via the MCPGODEBUG environment variable. +// +// The value of MCPGODEBUG is a comma-separated list of key=value pairs. +// For example: +// +// MCPGODEBUG=someoption=1,otheroption=value +package mcpgodebug + +import ( + "fmt" + "os" + "strings" +) + +const compatibilityEnvKey = "MCPGODEBUG" + +var compatibilityParams map[string]string + +func init() { + var err error + compatibilityParams, err = parseCompatibility(os.Getenv(compatibilityEnvKey)) + if err != nil { + panic(err) + } +} + +// Value returns the value of the compatibility parameter with the given key. +// It returns an empty string if the key is not set. +func Value(key string) string { + return compatibilityParams[key] +} + +func parseCompatibility(envValue string) (map[string]string, error) { + if envValue == "" { + return nil, nil + } + + params := make(map[string]string) + for part := range strings.SplitSeq(envValue, ",") { + k, v, ok := strings.Cut(part, "=") + if !ok { + return nil, fmt.Errorf("MCPGODEBUG: invalid format: %q", part) + } + params[strings.TrimSpace(k)] = strings.TrimSpace(v) + } + return params, nil +} diff --git a/internal/mcpgodebug/mcpgodebug_test.go b/internal/mcpgodebug/mcpgodebug_test.go new file mode 100644 index 0000000..9bbbc32 --- /dev/null +++ b/internal/mcpgodebug/mcpgodebug_test.go @@ -0,0 +1,88 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package mcpgodebug + +import ( + "testing" + + "github.com/google/go-cmp/cmp" +) + +func TestParseCompatibility_Success(t *testing.T) { + tests := []struct { + name string + envVal string + want map[string]string + }{ + { + name: "Basic", + envVal: "foo=bar,baz=qux", + want: map[string]string{ + "foo": "bar", + "baz": "qux", + }, + }, + { + name: "Empty", + envVal: "", + want: nil, + }, + { + name: "WithWhitespace", + envVal: " foo = bar \t, baz = qux ", + want: map[string]string{ + "foo": "bar", + "baz": "qux", + }, + }, + { + name: "WithEqualsSignInValue", + envVal: "foo=bar=baz", + want: map[string]string{ + "foo": "bar=baz", + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := parseCompatibility(tt.envVal) + if err != nil { + t.Fatalf("parseCompatibility() failed: %v", err) + } + if diff := cmp.Diff(tt.want, got); diff != "" { + t.Errorf("parseCompatibility() mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestParseCompatibility_Failure(t *testing.T) { + tests := []struct { + name string + envVal string + }{ + { + name: "NoEqualsSign", + envVal: "invalidformat", + }, + { + name: "MixedValidAndInvalid", + envVal: "foo=bar,baz", + }, + { + name: "EmptyPart", + envVal: "foo=bar,,baz=qux", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := parseCompatibility(tt.envVal) + if err == nil { + t.Error("parseCompatibility() expected error, got nil") + } + }) + } +} diff --git a/internal/oauthtest/fake_authorization_server.go b/internal/oauthtest/fake_authorization_server.go new file mode 100644 index 0000000..4cd6826 --- /dev/null +++ b/internal/oauthtest/fake_authorization_server.go @@ -0,0 +1,451 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package oauthtest + +import ( + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "maps" + "net/http" + "net/http/httptest" + "net/url" + "slices" + "testing" + + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" + "github.com/modelcontextprotocol/go-sdk/oauthex" +) + +type ClientInfo struct { + Secret string + RedirectURIs []string +} + +type MetadataEndpointConfig struct { + // Whether to serve the OAuth Authorization Server Metadata at + // /.well-known/oauth-authorization-server + issuerPath. + ServeOAuthInsertedEndpoint bool + // Whether to serve the OAuth Authorization Server Metadata at + // /.well-known/openid-configuration + issuerPath. + ServeOpenIDInsertedEndpoint bool + // Whether to serve the OAuth Authorization Server Metadata at + // issuerPath + /.well-known/openid-configuration. + // Should be used when issuerPath is not empty. + ServeOpenIDAppendedEndpoint bool +} + +type RegistrationConfig struct { + // Whether the client ID metadata document is supported. + ClientIDMetadataDocumentSupported bool + // PreregisteredClients is a map of valid ClientIDs to ClientSecrets. + PreregisteredClients map[string]ClientInfo + // Whether dynamic client registration is enabled. + DynamicClientRegistrationEnabled bool +} + +// JWTBearerConfig configures support for the JWT Bearer grant type (RFC 7523) +// on a [FakeAuthorizationServer]. +type JWTBearerConfig struct { + // ValidAssertions is the set of assertion values that are accepted. + // If empty, any non-empty assertion is accepted. + ValidAssertions []string +} + +// ClientCredentialsConfig configures support for the client_credentials +// grant type (RFC 6749 Section 4.4) on a [FakeAuthorizationServer]. +type ClientCredentialsConfig struct { + // Enabled controls whether the /token endpoint accepts + // grant_type=client_credentials and returns an access token + // if client authentication succeeds. + Enabled bool +} + +// Config holds configuration for FakeAuthorizationServer. +type Config struct { + // The optional path component of the issuer URL. + // If non-empty, it should start with a "/". It should not end with a "/". + // It affects the paths of the server endpoints. + IssuerPath string + // Configuration of the metadata endpoint. + MetadataEndpointConfig *MetadataEndpointConfig + // Configuration for client registration. + RegistrationConfig *RegistrationConfig + // JWTBearerConfig enables RFC 7523 JWT Bearer grant at the /token endpoint. + // If non-nil, the server accepts grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer. + JWTBearerConfig *JWTBearerConfig + // ClientCredentialsConfig enables RFC 6749 Section 4.4 client credentials + // grant at the /token endpoint. + ClientCredentialsConfig *ClientCredentialsConfig + // ScopesSupported is an optional list of scopes to advertise in the + // authorization server metadata. + ScopesSupported []string + // TokenScopeFunc, if set, is called with the scope from the authorization + // request and returns the scope string to include in the token response. + TokenScopeFunc func(requestedScope string) string + // AccessTokenTTL, if non-zero, is the expires_in (in seconds) returned by the + // /token endpoint for both the authorization_code and refresh_token grants. + // When zero a default of 3600 is used. Set it small to force a client's reuse + // token source to treat the access token as expired and refresh it. + AccessTokenTTL int + // IssueRefreshToken, if true, includes a refresh_token in token responses and + // enables grant_type=refresh_token at the /token endpoint. + IssueRefreshToken bool + // UnadvertiseIssParameter, if true, omits + // authorization_response_iss_parameter_supported from metadata while the + // authorize endpoint still includes iss. Used to test local policy for + // unadvertised RFC 9207 iss. + UnadvertiseIssParameter bool +} + +// testRefreshToken is the refresh token issued and accepted by the fake server +// when Config.IssueRefreshToken is set. +const testRefreshToken = "test_refresh_token" + +// accessTokenExpiresIn returns the expires_in value to use in token responses. +func (s *FakeAuthorizationServer) accessTokenExpiresIn() int { + if s.config.AccessTokenTTL != 0 { + return s.config.AccessTokenTTL + } + return 3600 +} + +// FakeAuthorizationServer is a fake OAuth 2.0 Authorization Server for testing. +type FakeAuthorizationServer struct { + server *httptest.Server + Mux *http.ServeMux + config Config + clients map[string]ClientInfo + codes map[string]codeInfo +} + +type codeInfo struct { + CodeChallenge string + Scope string +} + +// NewFakeAuthorizationServer creates a new FakeAuthorizationServer. +// The server is simple and should not be used outside of testing. +// It supports: +// - Only the authorization Code Grant +// - PKCE verification +// - Client tracking & dynamic registration +// - Client authentication +func NewFakeAuthorizationServer(config Config) *FakeAuthorizationServer { + s := &FakeAuthorizationServer{ + Mux: http.NewServeMux(), + config: config, + codes: make(map[string]codeInfo), + } + if config.RegistrationConfig != nil { + s.clients = maps.Clone(config.RegistrationConfig.PreregisteredClients) + } + if s.clients == nil { + s.clients = make(map[string]ClientInfo) + } + + s.Mux.HandleFunc(s.config.IssuerPath+"/authorize", s.handleAuthorize) + s.Mux.HandleFunc(s.config.IssuerPath+"/token", s.handleToken) + if config.MetadataEndpointConfig != nil { + if config.MetadataEndpointConfig.ServeOAuthInsertedEndpoint { + s.Mux.HandleFunc("/.well-known/oauth-authorization-server"+s.config.IssuerPath, s.handleMetadata) + } + if config.MetadataEndpointConfig.ServeOpenIDInsertedEndpoint { + s.Mux.HandleFunc("/.well-known/openid-configuration"+s.config.IssuerPath, s.handleMetadata) + } + if config.MetadataEndpointConfig.ServeOpenIDAppendedEndpoint && s.config.IssuerPath != "" { + s.Mux.HandleFunc(s.config.IssuerPath+"/.well-known/openid-configuration", s.handleMetadata) + } + } else { + // Serve the default OAuth endpoint. + s.Mux.HandleFunc("/.well-known/oauth-authorization-server", s.handleMetadata) + } + if config.RegistrationConfig != nil && config.RegistrationConfig.DynamicClientRegistrationEnabled { + s.Mux.HandleFunc(s.config.IssuerPath+"/register", s.handleRegister) + } + s.server = httptest.NewUnstartedServer(s.Mux) + + return s +} + +// Start starts the HTTP server and registers a cleanup function on t to close the server. +func (s *FakeAuthorizationServer) Start(t testing.TB) { + s.server.Start() + t.Cleanup(s.server.Close) +} + +// URL returns the base URL of the server (Issuer). +func (s *FakeAuthorizationServer) URL() string { + return s.server.URL +} + +func (s *FakeAuthorizationServer) handleMetadata(w http.ResponseWriter, r *http.Request) { + cimdSupported := false + var registrationEndpoint string + if s.config.RegistrationConfig != nil { + cimdSupported = s.config.RegistrationConfig.ClientIDMetadataDocumentSupported + if s.config.RegistrationConfig.DynamicClientRegistrationEnabled { + registrationEndpoint = s.URL() + s.config.IssuerPath + "/register" + } + } + meta := &oauthex.AuthServerMeta{ + Issuer: s.URL() + s.config.IssuerPath, + AuthorizationEndpoint: s.URL() + s.config.IssuerPath + "/authorize", + TokenEndpoint: s.URL() + s.config.IssuerPath + "/token", + RegistrationEndpoint: registrationEndpoint, + ScopesSupported: s.config.ScopesSupported, + ResponseTypesSupported: []string{"code"}, + CodeChallengeMethodsSupported: []string{"S256"}, + ClientIDMetadataDocumentSupported: cimdSupported, + TokenEndpointAuthMethodsSupported: []string{"client_secret_post", "client_secret_basic"}, + // Advertise RFC 9207 support unless UnadvertiseIssParameter is set: + // the authorize endpoint still includes "iss" in responses either way. + AuthorizationResponseIssParameterSupported: !s.config.UnadvertiseIssParameter, + } + // Set CORS headers for cross-origin client discovery. + w.Header().Set("Access-Control-Allow-Origin", "*") + w.Header().Set("Access-Control-Allow-Methods", "GET, OPTIONS") + w.Header().Set("Access-Control-Allow-Headers", "Content-Type") + + // Handle CORS preflight requests + if r.Method == http.MethodOptions { + w.WriteHeader(http.StatusNoContent) + return + } + + // Only GET allowed for metadata retrieval + if r.Method != http.MethodGet { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(meta); err != nil { + http.Error(w, "Failed to encode metadata", http.StatusInternalServerError) + return + } +} + +func (s *FakeAuthorizationServer) handleRegister(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + var metadata oauthex.ClientRegistrationMetadata + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, "failed to read request body", http.StatusBadRequest) + return + } + if err := internaljson.Unmarshal(body, &metadata); err != nil { + http.Error(w, "failed to parse request", http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusCreated) + clientID := rand.Text() + ci := ClientInfo{ + Secret: rand.Text(), + RedirectURIs: metadata.RedirectURIs, + } + s.clients[clientID] = ci + metadata.TokenEndpointAuthMethod = "client_secret_basic" + json.NewEncoder(w).Encode(&oauthex.ClientRegistrationResponse{ + ClientID: clientID, + ClientSecret: ci.Secret, + ClientRegistrationMetadata: metadata, + }) +} + +func (s *FakeAuthorizationServer) handleAuthorize(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + clientID := r.URL.Query().Get("client_id") + clientInfo, ok := s.clients[clientID] + if !ok { + http.Error(w, "unknown client_id", http.StatusBadRequest) + return + } + + redirectURI := r.URL.Query().Get("redirect_uri") + if redirectURI == "" { + http.Error(w, "missing redirect_uri", http.StatusBadRequest) + return + } + if !slices.Contains(clientInfo.RedirectURIs, redirectURI) { + http.Error(w, "invalid redirect_uri", http.StatusBadRequest) + return + } + codeChallenge := r.URL.Query().Get("code_challenge") + if codeChallenge == "" { + http.Error(w, "missing code_challenge", http.StatusBadRequest) + return + } + code := rand.Text() + s.codes[code] = codeInfo{ + CodeChallenge: codeChallenge, + Scope: r.URL.Query().Get("scope"), + } + + state := r.URL.Query().Get("state") + issuer := s.URL() + s.config.IssuerPath + + redirectURL := fmt.Sprintf("%s?code=%s&state=%s&iss=%s", redirectURI, code, state, url.QueryEscape(issuer)) + http.Redirect(w, r, redirectURL, http.StatusFound) +} + +func (s *FakeAuthorizationServer) handleToken(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + http.Error(w, "failed to parse form", http.StatusBadRequest) + return + } + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if err := s.authenticateClient(r); err != nil { + http.Error(w, err.Error(), http.StatusUnauthorized) + return + } + + grantType := r.Form.Get("grant_type") + switch grantType { + case "authorization_code": + s.handleAuthorizationCodeGrant(w, r) + case "urn:ietf:params:oauth:grant-type:jwt-bearer": + s.handleJWTBearerGrant(w, r) + case "client_credentials": + s.handleClientCredentialsGrant(w, r) + case "refresh_token": + s.handleRefreshTokenGrant(w, r) + default: + http.Error(w, fmt.Sprintf("unsupported grant_type: %s", grantType), http.StatusBadRequest) + } +} + +func (s *FakeAuthorizationServer) handleAuthorizationCodeGrant(w http.ResponseWriter, r *http.Request) { + code := r.Form.Get("code") + if code == "" { + http.Error(w, "missing code", http.StatusBadRequest) + return + } + codeInfo, ok := s.codes[code] + if !ok { + http.Error(w, "unknown authorization code", http.StatusBadRequest) + return + } + verifier := r.Form.Get("code_verifier") + if verifier == "" { + http.Error(w, "missing code_verifier", http.StatusBadRequest) + return + } + sha := sha256.Sum256([]byte(verifier)) + expectedChallenge := base64.RawURLEncoding.EncodeToString(sha[:]) + if expectedChallenge != codeInfo.CodeChallenge { + http.Error(w, "PKCE verification failed", http.StatusBadRequest) + return + } + + resp := map[string]any{ + "access_token": "test_access_token", + "token_type": "Bearer", + "expires_in": s.accessTokenExpiresIn(), + } + if s.config.IssueRefreshToken { + resp["refresh_token"] = testRefreshToken + } + if s.config.TokenScopeFunc != nil { + if scope := s.config.TokenScopeFunc(codeInfo.Scope); scope != "" { + resp["scope"] = scope + } + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) +} + +// handleRefreshTokenGrant implements grant_type=refresh_token (RFC 6749 Section +// 6) when Config.IssueRefreshToken is set, returning a distinct access token so +// callers can observe that a refresh occurred. +func (s *FakeAuthorizationServer) handleRefreshTokenGrant(w http.ResponseWriter, r *http.Request) { + if !s.config.IssueRefreshToken { + http.Error(w, "refresh_token grant not supported", http.StatusBadRequest) + return + } + if r.Form.Get("refresh_token") != testRefreshToken { + http.Error(w, "invalid refresh_token", http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "access_token": "test_access_token_refreshed", + "token_type": "Bearer", + "expires_in": s.accessTokenExpiresIn(), + "refresh_token": testRefreshToken, + }) +} + +func (s *FakeAuthorizationServer) handleJWTBearerGrant(w http.ResponseWriter, r *http.Request) { + if s.config.JWTBearerConfig == nil { + http.Error(w, "JWT bearer grant not supported", http.StatusBadRequest) + return + } + assertion := r.Form.Get("assertion") + if assertion == "" { + http.Error(w, "missing assertion", http.StatusBadRequest) + return + } + if len(s.config.JWTBearerConfig.ValidAssertions) > 0 { + if !slices.Contains(s.config.JWTBearerConfig.ValidAssertions, assertion) { + http.Error(w, "invalid assertion", http.StatusBadRequest) + return + } + } + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "access_token": "test_access_token", + "token_type": "Bearer", + "expires_in": 3600, + }) +} + +func (s *FakeAuthorizationServer) handleClientCredentialsGrant(w http.ResponseWriter, r *http.Request) { + if s.config.ClientCredentialsConfig == nil || !s.config.ClientCredentialsConfig.Enabled { + http.Error(w, "client_credentials grant not supported", http.StatusBadRequest) + return + } + resp := map[string]any{ + "access_token": "test_access_token", + "token_type": "Bearer", + "expires_in": 3600, + } + if s.config.TokenScopeFunc != nil { + if scope := s.config.TokenScopeFunc(r.Form.Get("scope")); scope != "" { + resp["scope"] = scope + } + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) +} + +func (s *FakeAuthorizationServer) authenticateClient(r *http.Request) error { + clientID, clientSecret, ok := r.BasicAuth() + if !ok { + clientID = r.Form.Get("client_id") + clientSecret = r.Form.Get("client_secret") + } + + clientInfo, ok := s.clients[clientID] + if !ok || clientInfo.Secret != clientSecret { + return errors.New("client not found") + } + return nil +} diff --git a/internal/oauthtest/fake_idp_server.go b/internal/oauthtest/fake_idp_server.go new file mode 100644 index 0000000..b7057ff --- /dev/null +++ b/internal/oauthtest/fake_idp_server.go @@ -0,0 +1,258 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package oauthtest + +import ( + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "fmt" + "maps" + "net/http" + "net/http/httptest" + "slices" + "testing" + "time" + + "github.com/modelcontextprotocol/go-sdk/oauthex" +) + +// TokenExchangeConfig configures RFC 8693 token exchange support on a [FakeIdPServer]. +type TokenExchangeConfig struct { + // IDJAGToken is the ID-JAG value returned from token exchange. + // Defaults to "test-id-jag-token" if empty. + IDJAGToken string +} + +// IdPConfig holds configuration for [FakeIdPServer]. +type IdPConfig struct { + // PreregisteredClients maps client IDs to their info. + PreregisteredClients map[string]ClientInfo + + // TokenExchangeConfig enables RFC 8693 token exchange at the /token endpoint. + // If non-nil, the server accepts grant_type=urn:ietf:params:oauth:grant-type:token-exchange. + TokenExchangeConfig *TokenExchangeConfig +} + +// FakeIdPServer is a fake OIDC Identity Provider for testing. +// It supports: +// - OIDC discovery (/.well-known/openid-configuration) +// - Authorization Code Grant with PKCE +// - ID Token issuance (fake JWTs) +// - RFC 8693 Token Exchange (ID Token → ID-JAG), if configured +type FakeIdPServer struct { + server *httptest.Server + mux *http.ServeMux + config IdPConfig + clients map[string]ClientInfo + codes map[string]codeInfo +} + +// NewFakeIdPServer creates a new FakeIdPServer. +func NewFakeIdPServer(config IdPConfig) *FakeIdPServer { + s := &FakeIdPServer{ + mux: http.NewServeMux(), + config: config, + clients: make(map[string]ClientInfo), + codes: make(map[string]codeInfo), + } + maps.Copy(s.clients, config.PreregisteredClients) + + s.mux.HandleFunc("/.well-known/openid-configuration", s.handleMetadata) + s.mux.HandleFunc("/authorize", s.handleAuthorize) + s.mux.HandleFunc("/token", s.handleToken) + s.server = httptest.NewUnstartedServer(s.mux) + + return s +} + +// Start starts the HTTP server and registers a cleanup function on t. +func (s *FakeIdPServer) Start(t testing.TB) { + s.server.Start() + t.Cleanup(s.server.Close) +} + +// URL returns the base URL of the server (issuer). +func (s *FakeIdPServer) URL() string { + return s.server.URL +} + +func (s *FakeIdPServer) handleMetadata(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + grantTypes := []string{"authorization_code"} + if s.config.TokenExchangeConfig != nil { + grantTypes = append(grantTypes, oauthex.GrantTypeTokenExchange) + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "issuer": s.URL(), + "authorization_endpoint": s.URL() + "/authorize", + "token_endpoint": s.URL() + "/token", + "jwks_uri": s.URL() + "/.well-known/jwks.json", + "response_types_supported": []string{"code"}, + "code_challenge_methods_supported": []string{"S256"}, + "grant_types_supported": grantTypes, + }) +} + +func (s *FakeIdPServer) handleAuthorize(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + clientID := r.URL.Query().Get("client_id") + clientInfo, ok := s.clients[clientID] + if !ok { + http.Error(w, "unknown client_id", http.StatusBadRequest) + return + } + redirectURI := r.URL.Query().Get("redirect_uri") + if redirectURI == "" { + http.Error(w, "missing redirect_uri", http.StatusBadRequest) + return + } + if !slices.Contains(clientInfo.RedirectURIs, redirectURI) { + http.Error(w, "invalid redirect_uri", http.StatusBadRequest) + return + } + codeChallenge := r.URL.Query().Get("code_challenge") + if codeChallenge == "" { + http.Error(w, "missing code_challenge", http.StatusBadRequest) + return + } + code := rand.Text() + s.codes[code] = codeInfo{CodeChallenge: codeChallenge} + state := r.URL.Query().Get("state") + redirectURL := fmt.Sprintf("%s?code=%s&state=%s", redirectURI, code, state) + http.Redirect(w, r, redirectURL, http.StatusFound) +} + +func (s *FakeIdPServer) handleToken(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if err := r.ParseForm(); err != nil { + http.Error(w, "failed to parse form", http.StatusBadRequest) + return + } + + if err := s.authenticateClient(r); err != nil { + http.Error(w, err.Error(), http.StatusUnauthorized) + return + } + + grantType := r.Form.Get("grant_type") + switch grantType { + case "authorization_code": + s.handleAuthorizationCodeGrant(w, r) + case oauthex.GrantTypeTokenExchange: + s.handleTokenExchangeGrant(w, r) + default: + http.Error(w, fmt.Sprintf("unsupported grant_type: %s", grantType), http.StatusBadRequest) + } +} + +func (s *FakeIdPServer) handleAuthorizationCodeGrant(w http.ResponseWriter, r *http.Request) { + code := r.Form.Get("code") + if code == "" { + http.Error(w, "missing code", http.StatusBadRequest) + return + } + ci, ok := s.codes[code] + if !ok { + http.Error(w, "unknown authorization code", http.StatusBadRequest) + return + } + verifier := r.Form.Get("code_verifier") + if verifier == "" { + http.Error(w, "missing code_verifier", http.StatusBadRequest) + return + } + sha := sha256.Sum256([]byte(verifier)) + expectedChallenge := base64.RawURLEncoding.EncodeToString(sha[:]) + if expectedChallenge != ci.CodeChallenge { + http.Error(w, "PKCE verification failed", http.StatusBadRequest) + return + } + + clientID := r.Form.Get("client_id") + now := time.Now().Unix() + idToken := fakeIDToken(s.URL(), clientID, now) + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "access_token": "test_access_token", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "test_refresh_token", + "id_token": idToken, + }) +} + +func (s *FakeIdPServer) handleTokenExchangeGrant(w http.ResponseWriter, r *http.Request) { + if s.config.TokenExchangeConfig == nil { + http.Error(w, "token exchange not supported", http.StatusBadRequest) + return + } + if r.Form.Get("requested_token_type") != oauthex.TokenTypeIDJAG { + http.Error(w, "invalid requested_token_type", http.StatusBadRequest) + return + } + if r.Form.Get("subject_token_type") != oauthex.TokenTypeIDToken { + http.Error(w, "invalid subject_token_type", http.StatusBadRequest) + return + } + if r.Form.Get("subject_token") == "" { + http.Error(w, "missing subject_token", http.StatusBadRequest) + return + } + + idJAG := s.config.TokenExchangeConfig.IDJAGToken + if idJAG == "" { + idJAG = "test-id-jag-token" + } + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "access_token": idJAG, + "issued_token_type": oauthex.TokenTypeIDJAG, + "token_type": "N_A", + "expires_in": 300, + }) +} + +func (s *FakeIdPServer) authenticateClient(r *http.Request) error { + clientID := r.Form.Get("client_id") + clientSecret := r.Form.Get("client_secret") + clientInfo, ok := s.clients[clientID] + if !ok { + return fmt.Errorf("unknown client") + } + if clientInfo.Secret != clientSecret { + return fmt.Errorf("invalid client credentials") + } + return nil +} + +// fakeIDToken creates a fake JWT ID token for testing. +// The token has a valid structure but is not cryptographically signed. +func fakeIDToken(issuer, audience string, now int64) string { + header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"RS256","typ":"JWT"}`)) + claimsJSON, _ := json.Marshal(map[string]any{ + "iss": issuer, + "sub": "test-user", + "aud": audience, + "exp": now + 3600, + "iat": now, + "email": "test@example.com", + }) + claims := base64.RawURLEncoding.EncodeToString(claimsJSON) + return header + "." + claims + ".mock-signature" +} diff --git a/internal/readme/README.src.md b/internal/readme/README.src.md index cdc2cd6..ce1f6fc 100644 --- a/internal/readme/README.src.md +++ b/internal/readme/README.src.md @@ -3,6 +3,7 @@ [![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/modelcontextprotocol/go-sdk) [![PkgGoDev](https://pkg.go.dev/badge/github.com/modelcontextprotocol/go-sdk)](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) +[![OpenSSF Scorecard](https://api.scorecard.dev/projects/github.com/modelcontextprotocol/go-sdk/badge)](https://scorecard.dev/viewer/?uri=github.com/modelcontextprotocol/go-sdk) This repository contains an implementation of the official Go software development kit (SDK) for the Model Context Protocol (MCP). @@ -28,6 +29,31 @@ The SDK consists of several importable packages: The SDK endeavors to implement the full MCP spec. The [`docs/`](/docs/) directory contains feature documentation, mapping the MCP spec to the packages above. +## Version Compatibility + +The following table shows which versions of the Go SDK support which versions of the MCP specification: + +| SDK Version | Latest MCP Spec | All Supported MCP Specs | +|-----------------|-------------------|------------------------------------------------------------------| +| v1.7.0+ | 2026-07-28 | 2026-07-28, 2025-11-25\*, 2025-06-18, 2025-03-26, 2024-11-05 | +| v1.4.0 - v1.6.1 | 2025-11-25\* | 2025-11-25\*, 2025-06-18, 2025-03-26, 2024-11-05 | +| v1.2.0 - v1.3.1 | 2025-11-25\*\* | 2025-11-25\*\*, 2025-06-18, 2025-03-26, 2024-11-05 | +| v1.0.0 - v1.1.0 | 2025-06-18 | 2025-06-18, 2025-03-26, 2024-11-05 | + +\* Client side OAuth has experimental support. + +\*\* Partial support for 2025-11-25 (client side OAuth and Sampling with tools not available). + +The roots, sampling, and logging features are deprecated as of protocol version +2026-07-28 by +[SEP-2577](https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging). +The SDK continues to support them for compatibility during the deprecation +window (at least twelve months). See the individual feature documentation for +migration guidance. + +New releases of the SDK target only supported versions of Go. See +https://go.dev/doc/devel/release#policy for more information. + ## Getting started To get started creating an MCP server, create an `mcp.Server` instance, add @@ -64,5 +90,5 @@ is a thriving ecosystem of Go MCP clients and servers. ## License -This project is licensed under the MIT License - see the [LICENSE](./LICENSE) -file for details. +This project is licensed under Apache 2.0 for new contributions, with existing +code under MIT - see the [LICENSE](./LICENSE) file for details. diff --git a/internal/readme/contributing.src.md b/internal/readme/contributing.src.md new file mode 100644 index 0000000..f2e931b --- /dev/null +++ b/internal/readme/contributing.src.md @@ -0,0 +1,273 @@ +# Contributing to the Go MCP SDK + +Thank you for your interest in contributing! The Go SDK needs active +contributions to keep up with changes in the MCP spec, fix bugs, and accommodate +new and emerging use-cases. We welcome all forms of contribution, from filing +and reviewing issues, to contributing fixes, to proposing and implementing new +features. + +As described in the [design document](https://github.com/modelcontextprotocol/go-sdk/blob/main/design/design.md), it is important for +the MCP SDK to remain idiomatic, future-proof, and extensible. The process +described here is intended to ensure that the SDK evolves safely and +transparently, while adhering to these goals. + +%toc + + +## Development setup + +This module can be built and tested using the standard Go toolchain. Run `go +test ./...` to run all its tests. + +To test changes to this module against another module that uses the SDK, we +recommend using a [`go.work` file](https://go.dev/doc/tutorial/workspaces) to +define a multi-module workspace. For example, if your directory contains a +`project` directory containing your project, and a `go-sdk` directory +containing the SDK, run: + +```sh +go work init ./project ./go-sdk +``` + +### Conformance tests + +The SDK includes a script to run the official MCP conformance tests against the +SDK's conformance server: + +```sh +./scripts/conformance.sh +``` + +By default, results are cleaned up after the script runs. To save results to a +specific directory: + +```sh +./scripts/conformance.sh --result_dir ./conformance-results +``` + +To run against a local checkout of the +[conformance repo](https://github.com/modelcontextprotocol/conformance) instead +of the latest npm release: + +```sh +./scripts/conformance.sh --conformance_repo ~/src/conformance +``` + +Note: you must run `npm install` in the conformance repo first. + +Run `./scripts/conformance.sh --help` for more options. + +## Filing issues + +This project uses the [GitHub issue +tracker](https://github.com/modelcontextprotocol/go-sdk/issues) for issues. The +process for filing bugs and proposals is described below. + +TODO(rfindley): describe a process for asking general questions in the public +MCP discord server. + +### Bugs + +Please [report +bugs](https://github.com/modelcontextprotocol/go-sdk/issues/new). If the SDK is +not working as you expected, it is likely due to a bug or inadequate +documentation, and reporting an issue will help us address this shortcoming. + +When reporting a bug, make sure to answer these five questions: + +1. What did you do? +2. What did you see? +3. What did you expect to see? +4. What version of the Go MCP SDK are you using? +5. What version of Go are you using (`go version`)? + +### Proposals + +A proposal is an issue that proposes a new API for the SDK, or a change to the +signature or behavior of an existing API. Proposals should be labeled with the +'proposal' label, and require an explicit approval from a maintainer before +being accepted (indicated by the 'proposal-accepted' label). Proposals must +remain open for at least a week to allow discussion before being accepted or +declined by a maintainer. + +Proposals that are straightforward and uncontroversial may be approved based on +discussion on the issue tracker or in a GitHub Discussion. However, proposals +that are deemed to be sufficiently unclear or complicated may be deferred to a +regular Working Group meeting (see 'Governance' below). + +This process is similar to the [Go proposal +process](https://github.com/golang/proposal), but is necessarily lighter weight +to accommodate the greater rate of change expected for the SDK. + +### Design discussion + +For open ended design discussion (anything that doesn't fall into the issue +categories above), use [GitHub +Discussions](https://github.com/modelcontextprotocol/go-sdk/discussions). +Ideally, each discussion should be focused on one aspect of the design. For +example: Tool Binding and Session APIs would be two separate discussions. +When discussions reach a consensus, they should be promoted into proposals. + +## Contributing code + +The project uses GitHub pull requests (PRs) to review changes. + +Any significant change should be associated with a GitHub issue. Issues that +are deemed to be good opportunities for contribution are be labeled ['Help +Wanted'](https://github.com/modelcontextprotocol/go-sdk/issues?q=is%3Aissue%20state%3Aopen%20label%3A%22help%20wanted%22). +If you want to work on such an issue, please first comment on the issue to say +that you're interested in contributing. For issues _not_ labeled 'Help Wanted', +it is recommended that you ask (and wait for confirmation) on the issue before +contributing, to avoid duplication of effort or wasted work. For nontrivial +changes that _don't_ relate to an existing issue, please file an issue first. + +Changes should be high quality and well tested, and should generally follow the +[Google Go style guide](https://google.github.io/styleguide/go/). Commit +messages should follow the [format used by the Go +project](https://go.dev/wiki/CommitMessage). + +Unless otherwise noted, the Go source files are distributed under the license +found in the LICENSE file. New contributions are licensed under Apache 2.0. All +Go files in the SDK should have a copyright header following the format below: + +```go +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. +``` + +### Adding and updating dependencies + +In general, the SDK tries to use as few dependencies as possible. Each new +dependency is a potential source for bugs, churn, and conflicts for our users. +Therefore, we require a [proposal](#proposals) for any new module dependency, +including upgrading an existing module to a new major version. New dependencies +should be evaluated for their stability and security, and should be +well-established in the Go ecosystem. + +In general, dependencies should be for internal use by the SDK implementation, +or for testing. Do not include types from dependencies in the SDK API. + +On the other hand, updating existing dependencies can be done at any time +without a proposal, as long as their major version does not change. Prefer to +update dependencies immediately following a release of the SDK, to allow as +much time as possible to find issues with the new version. + +After any change to dependencies, run govulncheck to check them for +vulnerabilities. + +``` +go run golang.org/x/vuln/cmd/govulncheck@latest +``` + +### Updating the README + +The top-level `README.md` file is generated from `internal/readme/README.src.md` +and should not be edited directly. To update the README: + +1. Make your changes to `internal/readme/README.src.md` +2. Run `go generate ./internal/readme` from the repository root to regenerate `README.md` +3. Commit both files together + +The CI system will automatically check that the README is up-to-date by running +`go generate ./internal/readme` and verifying no changes result. If you see a CI failure about the +README being out of sync, follow the steps above to regenerate it. + +## Versioning + +The SDK follows [semantic versioning](https://semver.org), as the Go module +system requires: the major version is part of the module path, so a breaking +release would import as `github.com/modelcontextprotocol/go-sdk/v2`. A breaking +change therefore cannot reach existing users by accident; they have to change +their import path to receive one. + +This policy covers the exported API of the SDK's importable packages — `mcp`, +`jsonrpc`, `auth`, `auth/extauth` and `oauthex`. Everything under `internal/` +is not importable outside the module and may change in any release. + +Which MCP spec revisions each SDK version speaks is documented in the +[README](README.md#version-compatibility). Support for a new revision arrives +in a minor release; dropping a revision the SDK previously negotiated is a +breaking change. + +### What counts as a breaking change + +- Removing or renaming an exported identifier. +- Changing the signature of an exported function or method. +- Removing or renaming a field of an exported struct, or changing its type. +- Adding a method to an exported interface that users are meant to implement. +- Changing documented behavior that callers could reasonably rely on. +- Dropping support for an MCP spec revision. + +The following are not breaking, and may land in a minor or patch release: + +- Adding an exported identifier, or a field to an exported struct. +- Fixing behavior that contradicted its documentation. +- Refactoring that leaves the exported API unchanged. +- Adding support for a new MCP spec revision. +- Raising the minimum Go version to one still supported upstream, per the + [README](README.md#version-compatibility). + +### Deprecation + +An API on its way out is marked with a `// Deprecated:` comment naming its +replacement, which `gopls` and `staticcheck` surface at call sites. It keeps +working until the next major version, since removing it would itself be a +breaking change. + +Protocol features that the *specification* deprecates are a separate matter and +follow the spec's timeline rather than this policy. The features deprecated by +[SEP-2577](https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging) +— roots, sampling and logging — remain supported for at least twelve months, as +noted in the [README](README.md#version-compatibility). + +Any change to the exported API, breaking or not, goes through the +[proposal process](#proposals). + +## Timeouts + +If a contributor hasn't responded to issue questions or PR comments in two weeks, +the issue or PR may be closed. It can be reopened when the contributor can resume +work. + +## Code of conduct + +This project follows the [Go Community Code of Conduct](https://go.dev/conduct). +If you encounter a conduct-related issue, please mail conduct@golang.org. + +## Governance + +Initially, the Go SDK repository will be administered by the Go team and +Anthropic, and they will be the approvers (the set of people able to merge PRs +to the SDK), also referred to as the 'Working Group'. The policies here are +also intended to satisfy necessary constraints of the Go team's participation +in the project. This may change in the future: see 'Ongoing Evaluation' below. + +### Working Group meetings + +On a regular basis, the Working Group will host a virtual meeting to discuss +outstanding proposals and other changes to the SDK. These meetings and their +agendas will be announced in advance, and open to all. The meetings will be +recorded, and recordings and meeting notes will be made available afterward. +(TODO: decide on a mechanism for tracking these meetings--likely a GitHub +issue.) + +This process is similar to the [Go Tools +call](https://go.dev/wiki/golang-tools), though it is expected that meetings +will at least initially occur on a more frequent basis. + +### Discord + +Discord (either the public or private Anthropic discord servers) should only be +used for logistical coordination or answering questions. For transparency and +durability, design discussion and decisions should occur in GitHub issues, +GitHub discussions, or public steering meetings. + +### Antitrust considerations + +The goal of this repository is to provide a robust and complete Go +implementation of the model context protocol, in an open and transparent +manner, without bias toward specific integration paths or providers. To that +end, the model context protocol organization's [antitrust +policy](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/ANTITRUST.md) +applies to all participation in this project. diff --git a/internal/readme/doc.go b/internal/readme/doc.go index 34bc60c..4461b94 100644 --- a/internal/readme/doc.go +++ b/internal/readme/doc.go @@ -2,7 +2,9 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. -//go:generate go run golang.org/x/example/internal/cmd/weave@latest -o ../../README.md ./README.src.md +//go:generate -command weave go run golang.org/x/example/internal/cmd/weave@latest +//go:generate weave -o ../../README.md ./README.src.md +//go:generate weave -o ../../CONTRIBUTING.md ./contributing.src.md // The readme package is used to generate README.md at the top-level of this // repo. Regenerate the README with go generate. diff --git a/internal/util/errors.go b/internal/util/errors.go new file mode 100644 index 0000000..ebeb614 --- /dev/null +++ b/internal/util/errors.go @@ -0,0 +1,16 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package util + +import ( + "fmt" +) + +// Wrapf wraps *errp with the given formatted message if *errp is not nil. +func Wrapf(errp *error, format string, args ...any) { + if *errp != nil { + *errp = fmt.Errorf("%s: %w", fmt.Sprintf(format, args...), *errp) + } +} diff --git a/internal/util/net.go b/internal/util/net.go new file mode 100644 index 0000000..593012d --- /dev/null +++ b/internal/util/net.go @@ -0,0 +1,49 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. +package util + +import ( + "net" + "net/netip" + "strings" +) + +func IsLoopback(addr string) bool { + host, _, err := net.SplitHostPort(addr) + if err != nil { + // If SplitHostPort fails, it might be just a host without a port. + host = strings.Trim(addr, "[]") + } + if host == "localhost" { + return true + } + ip, err := netip.ParseAddr(host) + if err != nil { + return false + } + return ip.IsLoopback() +} + +// cgnat is the RFC 6598 carrier-grade NAT range (100.64.0.0/10). netip's +// IsPrivate does not report it, but it is shared ISP space that is not publicly +// routable and could point at another tenant behind the same NAT. +var cgnat = netip.MustParsePrefix("100.64.0.0/10") + +// IsPrivateOrReserved reports whether ip belongs to a range that must not be +// reachable through a server-controlled URL, as a defense against SSRF: private +// (RFC 1918 / RFC 4193 ULA), link-local (including the 169.254.169.254 cloud +// metadata endpoint), carrier-grade NAT, multicast, and the unspecified address. +func IsPrivateOrReserved(ip netip.Addr) bool { + ip = ip.Unmap() + if !ip.IsValid() { + return true + } + if ip.IsPrivate() || + ip.IsLinkLocalUnicast() || + ip.IsMulticast() || + ip.IsUnspecified() { + return true + } + return cgnat.Contains(ip) +} diff --git a/internal/util/net_test.go b/internal/util/net_test.go new file mode 100644 index 0000000..21b7179 --- /dev/null +++ b/internal/util/net_test.go @@ -0,0 +1,87 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. +package util + +import ( + "net/netip" + "testing" +) + +// TestIsLoopback tests the IsLoopback helper function. +func TestIsLoopback(t *testing.T) { + tests := []struct { + addr string + want bool + }{ + {"localhost", true}, + {"localhost:3000", true}, + {"127.0.0.1", true}, + {"127.0.0.1:3000", true}, + {"[::1]", true}, + {"[::1]:3000", true}, + {"::1", true}, + {"", false}, + {"evil.com", false}, + {"evil.com:80", false}, + {"localhost.evil.com", false}, + {"127.0.0.1.evil.com", false}, + } + + for _, tt := range tests { + t.Run(tt.addr, func(t *testing.T) { + if got := IsLoopback(tt.addr); got != tt.want { + t.Errorf("IsLoopback(%q) = %v, want %v", tt.addr, got, tt.want) + } + }) + } +} + +func TestIsPrivateOrReserved(t *testing.T) { + tests := []struct { + ip string + want bool + }{ + // Private (RFC 1918) and IPv6 ULA (RFC 4193). + {"10.0.0.1", true}, + {"172.16.5.4", true}, + {"192.168.1.1", true}, + {"fd00::1", true}, + // Link-local, including the cloud metadata endpoint. + {"169.254.169.254", true}, + {"169.254.0.1", true}, + {"fe80::1", true}, + // Carrier-grade NAT (RFC 6598). + {"100.64.0.1", true}, + {"100.127.255.255", true}, + // Multicast and unspecified. + {"224.0.0.1", true}, + {"ff02::1", true}, + {"0.0.0.0", true}, + {"::", true}, + // IPv4-mapped IPv6 of a private address must also be caught. + {"::ffff:10.0.0.1", true}, + {"::ffff:192.168.0.1", true}, + // Loopback is intentionally NOT reserved here (policy handled elsewhere). + {"127.0.0.1", false}, + {"::1", false}, + // Public addresses are allowed. + {"8.8.8.8", false}, + {"1.1.1.1", false}, + {"2606:4700:4700::1111", false}, + // Just outside CGNAT. + {"100.63.255.255", false}, + {"100.128.0.0", false}, + } + for _, tt := range tests { + t.Run(tt.ip, func(t *testing.T) { + ip, err := netip.ParseAddr(tt.ip) + if err != nil { + t.Fatalf("ParseAddr(%q): %v", tt.ip, err) + } + if got := IsPrivateOrReserved(ip); got != tt.want { + t.Errorf("IsPrivateOrReserved(%q) = %v, want %v", tt.ip, got, tt.want) + } + }) + } +} diff --git a/internal/util/util.go b/internal/util/util.go deleted file mode 100644 index 4b5c325..0000000 --- a/internal/util/util.go +++ /dev/null @@ -1,44 +0,0 @@ -// Copyright 2025 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -package util - -import ( - "cmp" - "fmt" - "iter" - "slices" -) - -// Helpers below are copied from gopls' moremaps package. - -// Sorted returns an iterator over the entries of m in key order. -func Sorted[M ~map[K]V, K cmp.Ordered, V any](m M) iter.Seq2[K, V] { - // TODO(adonovan): use maps.Sorted if proposal #68598 is accepted. - return func(yield func(K, V) bool) { - keys := KeySlice(m) - slices.Sort(keys) - for _, k := range keys { - if !yield(k, m[k]) { - break - } - } - } -} - -// KeySlice returns the keys of the map M, like slices.Collect(maps.Keys(m)). -func KeySlice[M ~map[K]V, K comparable, V any](m M) []K { - r := make([]K, 0, len(m)) - for k := range m { - r = append(r, k) - } - return r -} - -// Wrapf wraps *errp with the given formatted message if *errp is not nil. -func Wrapf(errp *error, format string, args ...any) { - if *errp != nil { - *errp = fmt.Errorf("%s: %w", fmt.Sprintf(format, args...), *errp) - } -} diff --git a/jsonrpc/jsonrpc.go b/jsonrpc/jsonrpc.go index 1633d4e..a9ea78f 100644 --- a/jsonrpc/jsonrpc.go +++ b/jsonrpc/jsonrpc.go @@ -17,6 +17,8 @@ type ( Request = jsonrpc2.Request // Response is a JSON-RPC response. Response = jsonrpc2.Response + // Error is a structured error in a JSON-RPC response. + Error = jsonrpc2.WireError ) // MakeID coerces the given Go value to an ID. The value should be the @@ -37,3 +39,18 @@ func EncodeMessage(msg Message) ([]byte, error) { func DecodeMessage(data []byte) (Message, error) { return jsonrpc2.DecodeMessage(data) } + +// Standard JSON-RPC 2.0 error codes. +// See https://www.jsonrpc.org/specification#error_object +const ( + // CodeParseError indicates invalid JSON was received by the server. + CodeParseError = -32700 + // CodeInvalidRequest indicates the JSON sent is not a valid Request object. + CodeInvalidRequest = -32600 + // CodeMethodNotFound indicates the method does not exist or is not available. + CodeMethodNotFound = -32601 + // CodeInvalidParams indicates invalid method parameter(s). + CodeInvalidParams = -32602 + // CodeInternalError indicates an internal JSON-RPC error. + CodeInternalError = -32603 +) diff --git a/mcp/cache.go b/mcp/cache.go new file mode 100644 index 0000000..9837337 --- /dev/null +++ b/mcp/cache.go @@ -0,0 +1,87 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package mcp + +import ( + "sync" + "time" +) + +// methodCache is a per-method TTL cache for list and read results, as +// described in SEP-2549. Each entry is keyed by cursor (for paginated list +// methods) or URI (for resources/read). +type methodCache[R CacheableResult] struct { + mu sync.Mutex + cachedValues map[string]*cacheEntry[R] +} + +type cacheEntry[R CacheableResult] struct { + result R + receivedAt time.Time +} + +func (e *cacheEntry[R]) isValid() bool { + return time.Since(e.receivedAt) < time.Duration(e.result.GetTTLMs())*time.Millisecond +} + +func (mc *methodCache[R]) get(key string) (R, bool) { + mc.mu.Lock() + defer mc.mu.Unlock() + entry, ok := mc.cachedValues[key] + if !ok { + var zero R + return zero, false + } + if entry.result.GetTTLMs() <= 0 || !entry.isValid() { + delete(mc.cachedValues, key) + var zero R + return zero, false + } + return entry.result, true +} + +func (mc *methodCache[R]) put(key string, result R) { + mc.mu.Lock() + defer mc.mu.Unlock() + if mc.cachedValues == nil { + mc.cachedValues = make(map[string]*cacheEntry[R]) + } + mc.cachedValues[key] = &cacheEntry[R]{ + result: result, + receivedAt: time.Now(), + } +} + +func (mc *methodCache[R]) invalidate() { + mc.mu.Lock() + defer mc.mu.Unlock() + clear(mc.cachedValues) +} + +func (mc *methodCache[R]) invalidateKey(key string) { + mc.mu.Lock() + defer mc.mu.Unlock() + delete(mc.cachedValues, key) +} + +// cursorParams is the constraint for list-method params that carry a pagination +// cursor and can be checked for nil. Both methods are already implemented by +// every concrete list-params type. +type cursorParams interface { + Params + cursorPtr() *string +} + +// cachedListResult returns a cached list result keyed by the request cursor +// (SEP-2549). It returns the zero value and false on miss or when params is nil. +func cachedListResult[P cursorParams, R CacheableResult](cache *methodCache[R], params P) (R, bool) { + key := "" + if !params.isNil() { + if cp := params.cursorPtr(); cp != nil { + key = *cp + } + } + return cache.get(key) +} diff --git a/mcp/capabilities_test.go b/mcp/capabilities_test.go new file mode 100644 index 0000000..86c9970 --- /dev/null +++ b/mcp/capabilities_test.go @@ -0,0 +1,184 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +//lint:file-ignore SA1019 tests exercise deprecated SEP-2577 APIs (roots, sampling, logging). + +package mcp + +import ( + "context" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/google/jsonschema-go/jsonschema" +) + +// TestServerListChangedNotifications verifies that listChanged notifications +// are correctly sent or suppressed based on capability configuration. +func TestServerListChangedNotifications(t *testing.T) { + tool := &Tool{Name: "test-tool", InputSchema: &jsonschema.Schema{Type: "object"}} + + testCases := []struct { + name string + serverOpts *ServerOptions + wantNotifyCount int64 + }{ + { + name: "Default: notification sent", + serverOpts: &ServerOptions{ + Capabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: true}, + }, + }, + wantNotifyCount: 1, + }, + { + name: "ListChanged false: notification suppressed", + serverOpts: &ServerOptions{ + Capabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: false}, + }, + }, + wantNotifyCount: 0, + }, + { + name: "ListChanged true: notification sent", + serverOpts: &ServerOptions{ + Capabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: true}, + }, + }, + wantNotifyCount: 1, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + + // Create server. + impl := &Implementation{Name: "testServer", Version: "v1.0.0"} + server := NewServer(impl, tc.serverOpts) + + // Track notifications. + var notifyCount atomic.Int64 + + // Connect client and server. + cTransport, sTransport := NewInMemoryTransports() + ss, err := server.Connect(ctx, sTransport, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, &ClientOptions{ + ToolListChangedHandler: func(ctx context.Context, req *ToolListChangedRequest) { + notifyCount.Add(1) + }, + }) + cs, err := client.Connect(ctx, cTransport, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Add a tool, which may or may not trigger notification. + server.AddTool(tool, nil) + + // Sleep an arbitrary time longer than the debounce delay (synctest + // makes this practical). + time.Sleep(1 * time.Second) + + // Wait for all goroutines to be blocked (notification delivered). + synctest.Wait() + + if got, want := notifyCount.Load(), tc.wantNotifyCount; got != want { + t.Errorf("notification count: got %d, want %d", got, want) + } + }) + }) + } +} + +// TestClientListChangedNotifications verifies that roots listChanged notifications +// are correctly sent or suppressed based on client capability configuration. +func TestClientListChangedNotifications(t *testing.T) { + root := &Root{URI: "file:///test"} + + testCases := []struct { + name string + clientOpts *ClientOptions + wantNotifyCount int64 + }{ + { + name: "Default: notification sent", + clientOpts: nil, + wantNotifyCount: 1, + }, + { + name: "ListChanged false: notification suppressed", + clientOpts: &ClientOptions{ + Capabilities: &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: false}, + }, + }, + wantNotifyCount: 0, + }, + { + name: "ListChanged true: notification sent", + clientOpts: &ClientOptions{ + Capabilities: &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: true}, + }, + }, + wantNotifyCount: 1, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + + // Track notifications. + var notifyCount atomic.Int64 + + // Create server with roots list changed handler. + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, &ServerOptions{ + RootsListChangedHandler: func(ctx context.Context, req *RootsListChangedRequest) { + notifyCount.Add(1) + }, + }) + + // Connect client and server. + cTransport, sTransport := NewInMemoryTransports() + ss, err := server.Connect(ctx, sTransport, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, tc.clientOpts) + cs, err := client.Connect(ctx, cTransport, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Add a root, which may or may not trigger notification. + client.AddRoots(root) + + // Wait for all goroutines to be blocked (notification delivered). + synctest.Wait() + + if got, want := notifyCount.Load(), tc.wantNotifyCount; got != want { + t.Errorf("notification count: got %d, want %d", got, want) + } + }) + }) + } +} diff --git a/mcp/client.go b/mcp/client.go index d7e3ae5..1ee13e8 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -6,15 +6,20 @@ package mcp import ( "context" - "encoding/json" + "errors" "fmt" "iter" + "log/slog" + "maps" + "reflect" "slices" + "strings" "sync" "sync/atomic" "time" "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/internal/json" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) @@ -29,6 +34,11 @@ type Client struct { sessions []*ClientSession sendingMethodHandler_ MethodHandler receivingMethodHandler_ MethodHandler + // sendMethods is the list of methods this client may send to a + // server: it always contains the standard server methods (from + // serverMethodInfos) plus any custom methods registered via + // [AddSendingCustomMethod]. + sendMethods map[string]methodInfo } // NewClient creates a new [Client]. @@ -38,47 +48,170 @@ type Client struct { // The first argument must not be nil. // // If non-nil, the provided options configure the Client. -func NewClient(impl *Implementation, opts *ClientOptions) *Client { +func NewClient(impl *Implementation, options *ClientOptions) *Client { if impl == nil { panic("nil Implementation") } + var opts ClientOptions + if options != nil { + opts = *options + } + options = nil // prevent reuse + + if opts.CreateMessageHandler != nil && opts.CreateMessageWithToolsHandler != nil { + panic("cannot set both CreateMessageHandler and CreateMessageWithToolsHandler; use CreateMessageWithToolsHandler for tool support, or CreateMessageHandler for basic sampling") + } + if opts.Logger == nil { // ensure we have a logger + opts.Logger = ensureLogger(nil) + } + + sendMethods := make(map[string]methodInfo, len(serverMethodInfos)) + maps.Copy(sendMethods, serverMethodInfos) + c := &Client{ impl: impl, + opts: opts, roots: newFeatureSet(func(r *Root) string { return r.URI }), - sendingMethodHandler_: defaultSendingMethodHandler[*ClientSession], + sendingMethodHandler_: defaultSendingMethodHandler, receivingMethodHandler_: defaultReceivingMethodHandler[*ClientSession], + sendMethods: sendMethods, } - if opts != nil { - c.opts = *opts + if opts.MultiRoundTrip == nil || !opts.MultiRoundTrip.Disabled { + c.AddSendingMiddleware(clientMultiRoundTripMiddleware()) } return c } // ClientOptions configures the behavior of the client. type ClientOptions struct { + // Logger may be set to a non-nil value to enable logging of client activity. + Logger *slog.Logger // CreateMessageHandler handles incoming requests for sampling/createMessage. // - // Setting CreateMessageHandler to a non-nil value causes the client to - // advertise the sampling capability. + // Setting CreateMessageHandler to a non-nil value automatically causes the + // client to advertise the sampling capability, with default value + // &SamplingCapabilities{}. If [ClientOptions.Capabilities] is set and has a + // non nil value for [ClientCapabilities.Sampling], that value overrides the + // inferred capability. + // + // Deprecated: the sampling feature is deprecated as of protocol version + // 2026-07-28 (SEP-2577). It remains functional during the deprecation + // window (at least twelve months). Migrate to calling LLM provider APIs + // directly from your server. See + // https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. CreateMessageHandler func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) + // CreateMessageWithToolsHandler handles incoming sampling/createMessage + // requests that may involve tool use. It returns + // [CreateMessageWithToolsResult], which supports array content for parallel + // tool calls. + // + // Setting this handler causes the client to advertise the sampling + // capability with tools support (sampling.tools). As with + // [CreateMessageHandler], [ClientOptions.Capabilities].Sampling overrides + // the inferred capability. + // + // It is a panic to set both CreateMessageHandler and + // CreateMessageWithToolsHandler. + // + // Deprecated: the sampling feature is deprecated as of protocol version + // 2026-07-28 (SEP-2577). It remains functional during the deprecation + // window (at least twelve months). Migrate to calling LLM provider APIs + // directly from your server. See + // https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. + CreateMessageWithToolsHandler func(context.Context, *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) // ElicitationHandler handles incoming requests for elicitation/create. // - // Setting ElicitationHandler to a non-nil value causes the client to - // advertise the elicitation capability. + // Setting ElicitationHandler to a non-nil value automatically causes the + // client to advertise the elicitation capability, with default value + // &ElicitationCapabilities{}. If [ClientOptions.Capabilities] is set and has + // a non nil value for [ClientCapabilities.ELicitattion], that value + // overrides the inferred capability. ElicitationHandler func(context.Context, *ElicitRequest) (*ElicitResult, error) + // Capabilities optionally configures the client's default capabilities, + // before any capabilities are inferred from other configuration. + // + // If Capabilities is nil, the default client capabilities are + // {"roots":{"listChanged":true}}, for historical reasons. Setting + // Capabilities to a non-nil value overrides this default. As a special case, + // to work around #607, Capabilities.Roots is ignored: set + // Capabilities.RootsV2 to configure the roots capability. This allows the + // "roots" capability to be disabled entirely. + // + // For example: + // - To disable the "roots" capability, use &ClientCapabilities{} + // - To configure "roots", but disable "listChanged" notifications, use + // &ClientCapabilities{RootsV2:&RootCapabilities{}}. + // + // # Interaction with capability inference + // + // Sampling and elicitation capabilities are automatically added when their + // corresponding handlers are set, with the default value described at + // [ClientOptions.CreateMessageHandler] and + // [ClientOptions.ElicitationHandler]. If the Sampling or Elicitation fields + // are set in the Capabilities field, their values override the inferred + // value. + // + // For example, to advertise sampling with tools and context support: + // + // Capabilities: &ClientCapabilities{ + // Sampling: &SamplingCapabilities{ + // Tools: &SamplingToolsCapabilities{}, + // Context: &SamplingContextCapabilities{}, + // }, + // } + // + // Or to configure elicitation modes: + // + // Capabilities: &ClientCapabilities{ + // Elicitation: &ElicitationCapabilities{ + // Form: &FormElicitationCapabilities{}, + // URL: &URLElicitationCapabilities{}, + // }, + // } + // + // Conversely, if Capabilities does not set a field (for example, if the + // Elicitation field is nil), the inferred capability will be used. + Capabilities *ClientCapabilities + // ElicitationCompleteHandler handles incoming notifications for notifications/elicitation/complete. + ElicitationCompleteHandler func(context.Context, *ElicitationCompleteNotificationRequest) // Handlers for notifications from the server. - ToolListChangedHandler func(context.Context, *ToolListChangedRequest) - PromptListChangedHandler func(context.Context, *PromptListChangedRequest) - ResourceListChangedHandler func(context.Context, *ResourceListChangedRequest) - ResourceUpdatedHandler func(context.Context, *ResourceUpdatedNotificationRequest) + ToolListChangedHandler func(context.Context, *ToolListChangedRequest) + PromptListChangedHandler func(context.Context, *PromptListChangedRequest) + ResourceListChangedHandler func(context.Context, *ResourceListChangedRequest) + ResourceUpdatedHandler func(context.Context, *ResourceUpdatedNotificationRequest) + // LoggingMessageHandler handles incoming notifications/message + // notifications from the server. + // + // Deprecated: the logging feature is deprecated as of protocol version + // 2026-07-28 (SEP-2577). It remains functional during the deprecation + // window (at least twelve months). See + // https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. LoggingMessageHandler func(context.Context, *LoggingMessageRequest) ProgressNotificationHandler func(context.Context, *ProgressNotificationClientRequest) + // MultiRoundTrip configures the automatic MultiRoundTrip (Multi Round-Trip Requests) middleware. + // By default (nil), the middleware is enabled with default settings. + // Set Disabled to true to opt out of automatic MultiRoundTrip handling. + MultiRoundTrip *MultiRoundTripOptions // If non-zero, defines an interval for regular "ping" requests. // If the peer fails to respond to pings originating from the keepalive check, // the session is automatically closed. + // NOTE: The keepalive feature is only available for protocol versions < 2026-07-28 KeepAlive time.Duration + // KeepAliveFailureThreshold is the number of consecutive keepalive ping + // failures tolerated before the session is closed. A value of 0 or 1 + // closes the session on the first failure (the default). Higher values + // align with the spec's "multiple failed pings MAY trigger a connection + // reset" guidance, letting a transient miss pass without tearing down an + // otherwise live session. Has no effect unless KeepAlive is non-zero. + KeepAliveFailureThreshold int } +// toolContextKeyType is the context key type for passing tool definitions +// from CallTool to the transport layer. +type toolContextKeyType struct{} + +var toolContextKey = toolContextKeyType{} + // bind implements the binder[*ClientSession] interface, so that Clients can // be connected using [connect]. func (c *Client) bind(mcpConn Connection, conn *jsonrpc2.Connection, state *clientSessionState, onClose func()) *ClientSession { @@ -112,17 +245,55 @@ func (e unsupportedProtocolVersionError) Error() string { return fmt.Sprintf("unsupported protocol version: %q", e.version) } -// ClientSessionOptions is reserved for future use. -type ClientSessionOptions struct{} +// ClientSessionOptions configures a client session created by [Client.Connect]. +type ClientSessionOptions struct { + // ProtocolVersion is the protocol version sent in the initialize (or + // discover) request. If empty, the latest supported version is used. + // The server may negotiate a different mutually supported version. + ProtocolVersion string +} -func (c *Client) capabilities() *ClientCapabilities { - caps := &ClientCapabilities{} - caps.Roots.ListChanged = true - if c.opts.CreateMessageHandler != nil { - caps.Sampling = &SamplingCapabilities{} +func (c *Client) capabilities(protocolVersion string) *ClientCapabilities { + // Start with user-provided capabilities as defaults, or use SDK defaults. + var caps *ClientCapabilities + if c.opts.Capabilities != nil { + // Deep copy the user-provided capabilities to avoid mutation. + caps = c.opts.Capabilities.clone() + } else { + // SDK defaults: roots with listChanged. + // (this was the default behavior at v1.0.0, and so cannot be changed) + caps = &ClientCapabilities{ + RootsV2: &RootCapabilities{ + ListChanged: true, + }, + } } + + // Sync Roots from RootsV2 for backward compatibility (#607). + if caps.RootsV2 != nil { + caps.Roots = *caps.RootsV2 + } + + // Augment with sampling capability if a handler is set. + if c.opts.CreateMessageHandler != nil || c.opts.CreateMessageWithToolsHandler != nil { + if caps.Sampling == nil { + caps.Sampling = &SamplingCapabilities{} + if c.opts.CreateMessageWithToolsHandler != nil { + caps.Sampling.Tools = &SamplingToolsCapabilities{} + } + } + } + + // Augment with elicitation capability if handler is set. if c.opts.ElicitationHandler != nil { - caps.Elicitation = &ElicitationCapabilities{} + if caps.Elicitation == nil { + caps.Elicitation = &ElicitationCapabilities{} + // Form elicitation was added in 2025-11-25; for older versions, + // {} is treated the same as {"form":{}}. + if protocolVersion >= protocolVersion20251125 { + caps.Elicitation.Form = &FormElicitationCapabilities{} + } + } } return caps } @@ -134,16 +305,86 @@ func (c *Client) capabilities() *ClientCapabilities { // when it is no longer needed. However, if the connection is closed by the // server, calls or notifications will return an error wrapping // [ErrConnectionClosed]. -func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptions) (cs *ClientSession, err error) { - cs, err = connect(ctx, t, c, (*clientSessionState)(nil), nil) +func (c *Client) Connect(ctx context.Context, t Transport, opts *ClientSessionOptions) (cs *ClientSession, err error) { + cs, err = connect(ctx, t, c, (*clientSessionState)(nil), nil, c.opts.Logger) if err != nil { return nil, err } + protocolVersion := latestProtocolVersion + if opts != nil && opts.ProtocolVersion != "" { + protocolVersion = opts.ProtocolVersion + } + + if protocolVersion >= protocolVersion20260728 { + // Per SEP-2575, try the stateless server/discover RPC first. If the server + // signals it doesn't support it, fall back to the legacy initialize + // handshake. + discoverCtx := context.WithValue(ctx, protocolVersionContextKey{}, protocolVersion) + // We try to discover the server's capabilities. If the server rejects the + // requested version but specifies which versions it supports, we negotiate + // a mutually supported version and try again. + for range 2 { + discRes, err := c.discover(discoverCtx, cs) + if err == nil { + cs.state.InitializeResult = discRes + if hc, ok := cs.mcpConn.(clientConnection); ok { + hc.sessionUpdated(cs.state) + } + subscribeParams := &SubscriptionsListenParams{ + Notifications: &NotificationSubscriptions{}, + } + caps := discRes.Capabilities + if caps == nil { + caps = &ServerCapabilities{} + } + if c.opts.ToolListChangedHandler != nil && caps.Tools != nil && caps.Tools.ListChanged { + subscribeParams.Notifications.ToolsListChanged = true + } + if c.opts.PromptListChangedHandler != nil && caps.Prompts != nil && caps.Prompts.ListChanged { + subscribeParams.Notifications.PromptsListChanged = true + } + if c.opts.ResourceListChangedHandler != nil && caps.Resources != nil && caps.Resources.ListChanged { + subscribeParams.Notifications.ResourcesListChanged = true + } + if subscribeParams.Notifications.ToolsListChanged || + subscribeParams.Notifications.PromptsListChanged || + subscribeParams.Notifications.ResourcesListChanged { + // ClientSession.Close cancels the listenCtx context to send notifications/cancelled. + listenCtx, cancelListen := context.WithCancel(context.Background()) + cs.listenCancel = cancelListen + if err := cs.subscriptionsListen(listenCtx, subscribeParams); err != nil { + cancelListen() + return nil, fmt.Errorf("opening subscriptions/listen: %w", err) + } + } + return cs, nil + } + + // Try to negotiate a mutually supported version if the server + // reports an UnsupportedProtocolVersionError with a supported version. + var werr *jsonrpc.Error + if errors.As(err, &werr) && werr.Code == CodeUnsupportedProtocolVersion && len(werr.Data) > 0 { + var data UnsupportedProtocolVersionData + if err := json.Unmarshal(werr.Data, &data); err == nil { + if negotiatedVersion := negotiateMutuallySupportedVersion(data.Supported); negotiatedVersion != "" && negotiatedVersion >= protocolVersion20260728 { + discoverCtx = context.WithValue(ctx, protocolVersionContextKey{}, negotiatedVersion) + continue + } + } + } + // Per the spec, fall back to the legacy initialize handshake on any + // non-modern error from server/discover. + break + } + // Use the latest legacy protocol version for the fallback initialize. + protocolVersion = protocolVersion20251125 + } + params := &InitializeParams{ - ProtocolVersion: latestProtocolVersion, + ProtocolVersion: protocolVersion, ClientInfo: c.impl, - Capabilities: c.capabilities(), + Capabilities: c.capabilities(protocolVersion), } req := &InitializeRequest{Session: cs, Params: params} res, err := handleSend[*InitializeResult](ctx, methodInitialize, req) @@ -152,6 +393,7 @@ func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptio return nil, err } if !slices.Contains(supportedProtocolVersions, res.ProtocolVersion) { + _ = cs.Close() return nil, unsupportedProtocolVersionError{res.ProtocolVersion} } cs.state.InitializeResult = res @@ -171,6 +413,54 @@ func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptio return cs, nil } +// discover sends a SEP-2575 server/discover request to probe the server for +// stateless protocol support. +func (c *Client) discover(ctx context.Context, cs *ClientSession) (*InitializeResult, error) { + protocolVersion := protocolVersionFromContext(ctx) + caps := c.capabilities(protocolVersion) + params := &DiscoverParams{ + Meta: Meta{ + MetaKeyProtocolVersion: protocolVersion, + MetaKeyClientInfo: c.impl, + MetaKeyClientCapabilities: caps.toV2(), + }, + } + req := &DiscoverRequest{Session: cs, Params: params} + res, err := handleSend[*DiscoverResult](ctx, methodDiscover, req) + if err != nil { + return nil, err + } + + // Pick the highest protocol version that both the server and this SDK support. + // Since supportedProtocolVersions is defined in descending order (newest to oldest), + // the first match we find is the highest supported version. + var negotiated string + if slices.Contains(res.SupportedVersions, protocolVersion) { + negotiated = protocolVersion + } else { + negotiated = negotiateMutuallySupportedVersion(res.SupportedVersions) + } + if negotiated == "" || negotiated < protocolVersion20260728 { + // If there is no overlap, fall back to initialize so version + // negotiation can happen via the legacy path. + return nil, &jsonrpc2.WireError{ + Code: CodeUnsupportedProtocolVersion, + Message: "unsupported protocol version", + } + } + + var serverInfo *Implementation + if v, ok := decodeMetaValue[*Implementation](res.GetMeta(), MetaKeyServerInfo); ok { + serverInfo = v + } + return &InitializeResult{ + Capabilities: res.Capabilities, + Instructions: res.Instructions, + ProtocolVersion: negotiated, + ServerInfo: serverInfo, + }, nil +} + // A ClientSession is a logical connection with an MCP server. Its // methods can be used to send requests or notifications to the server. Create // a session by calling [Client.Connect]. @@ -187,11 +477,32 @@ type ClientSession struct { conn *jsonrpc2.Connection client *Client keepaliveCancel context.CancelFunc + listenCancel context.CancelFunc mcpConn Connection // No mutex is (currently) required to guard the session state, because it is // only set synchronously during Client.Connect. state clientSessionState + + // Per-method TTL caches for list results (SEP-2549). + toolsCache methodCache[*ListToolsResult] + promptsCache methodCache[*ListPromptsResult] + resourcesCache methodCache[*ListResourcesResult] + resourceTemplatesCache methodCache[*ListResourceTemplatesResult] + readResourceCache methodCache[*ReadResourceResult] + + // Pending URL elicitations waiting for completion notifications. + pendingElicitationsMu sync.Mutex + pendingElicitations map[string]chan struct{} + + // resourceSubsMu guards resourceSubs. + resourceSubsMu sync.Mutex + // resourceSubs maps a subscribed resource URI to the cancel func of the + // goroutine running its dedicated subscriptions/listen stream. Populated + // only under SEP-2575; the legacy protocol routes Subscribe and + // Unsubscribe straight to the resources/subscribe and resources/unsubscribe + // RPCs and leaves this map untouched. + resourceSubs map[string]context.CancelFunc } type clientSessionState struct { @@ -200,6 +511,45 @@ type clientSessionState struct { func (cs *ClientSession) InitializeResult() *InitializeResult { return cs.state.InitializeResult } +// usesNewProtocol reports whether this session has negotiated a protocol +// version >= 2026-07-28, which requires the SEP-2575 per-request `_meta` +// triple on every outgoing request. +func (cs *ClientSession) usesNewProtocol() bool { + res := cs.state.InitializeResult + return res != nil && res.ProtocolVersion >= protocolVersion20260728 +} + +// injectRequestMeta populates the SEP-2575 per-request `_meta` fields +// (protocolVersion, optional clientInfo, clientCapabilities) on the given +// outgoing request params. Keys already present in params.Meta are not +// overwritten. Per PR modelcontextprotocol/modelcontextprotocol#3002 +// clientInfo is SHOULD (not MUST), and is omitted when the client has no +// [Implementation] configured. +func injectRequestMeta[T any, P interface { + *T + Params +}](cs *ClientSession, params P) P { + res := cs.state.InitializeResult + if params == nil { + params = new(T) + } + m := params.GetMeta() + if m == nil { + m = map[string]any{} + } + if _, ok := m[MetaKeyProtocolVersion]; !ok { + m[MetaKeyProtocolVersion] = res.ProtocolVersion + } + if _, ok := m[MetaKeyClientInfo]; !ok && cs.client.impl != nil { + m[MetaKeyClientInfo] = cs.client.impl + } + if _, ok := m[MetaKeyClientCapabilities]; !ok { + m[MetaKeyClientCapabilities] = cs.client.capabilities(res.ProtocolVersion).toV2() + } + params.SetMeta(m) + return params +} + func (cs *ClientSession) ID() string { if c, ok := cs.mcpConn.(hasSessionID); ok { return c.SessionID() @@ -214,13 +564,18 @@ func (cs *ClientSession) ID() string { // Close is idempotent and concurrency safe. func (cs *ClientSession) Close() error { // Note: keepaliveCancel access is safe without a mutex because: - // 1. keepaliveCancel is only written once during startKeepalive (happens-before all Close calls) + // 1. keepaliveCancel is only written once during Client.Connect (through startKeepalive), + // which happens before any code that may call Close from another goroutine // 2. context.CancelFunc is safe to call multiple times and from multiple goroutines // 3. The keepalive goroutine calls Close on ping failure, but this is safe since // Close is idempotent and conn.Close() handles concurrent calls correctly if cs.keepaliveCancel != nil { cs.keepaliveCancel() } + if cs.listenCancel != nil { + cs.listenCancel() + } + cs.cancelAllResourceSubscriptions() err := cs.conn.Close() if cs.onClose != nil && cs.calledOnClose.CompareAndSwap(false, true) { @@ -236,14 +591,78 @@ func (cs *ClientSession) Wait() error { return cs.conn.Wait() } +// lookupTool returns the most recently seen definition of the tool with the +// given name across all cached ListTools results, or nil if no such tool has +// been seen. It is used by CallTool to inject the tool definition into the +// outgoing request context for transport-layer features (e.g. x-mcp-header +// param annotations). +func (cs *ClientSession) lookupTool(name string) *Tool { + cs.toolsCache.mu.Lock() + defer cs.toolsCache.mu.Unlock() + for _, entry := range cs.toolsCache.cachedValues { + for _, t := range entry.result.Tools { + if t.Name == name { + return t + } + } + } + return nil +} + +// registerElicitationWaiter registers a waiter for an elicitation complete +// notification with the given elicitation ID. It returns two functions: an await +// function that waits for the notification or context cancellation, and a cleanup +// function that must be called to unregister the waiter. This must be called before +// triggering the elicitation to avoid a race condition where the notification +// arrives before the waiter is registered. +// +// The cleanup function must be called even if the await function is never called, +// to prevent leaking the registration. +func (cs *ClientSession) registerElicitationWaiter(elicitationID string) (await func(context.Context) error, cleanup func()) { + // Create a channel for this elicitation. + ch := make(chan struct{}, 1) + + // Register the channel. + cs.pendingElicitationsMu.Lock() + if cs.pendingElicitations == nil { + cs.pendingElicitations = make(map[string]chan struct{}) + } + cs.pendingElicitations[elicitationID] = ch + cs.pendingElicitationsMu.Unlock() + + // Return await and cleanup functions. + await = func(ctx context.Context) error { + select { + case <-ctx.Done(): + return fmt.Errorf("context cancelled while waiting for elicitation completion: %w", ctx.Err()) + case <-ch: + return nil + } + } + + cleanup = func() { + cs.pendingElicitationsMu.Lock() + delete(cs.pendingElicitations, elicitationID) + cs.pendingElicitationsMu.Unlock() + } + + return await, cleanup +} + // startKeepalive starts the keepalive mechanism for this client session. func (cs *ClientSession) startKeepalive(interval time.Duration) { - startKeepalive(cs, interval, &cs.keepaliveCancel) + startKeepalive(cs, interval, cs.client.opts.KeepAliveFailureThreshold, &cs.keepaliveCancel, cs.client.opts.Logger) } // AddRoots adds the given roots to the client, // replacing any with the same URIs, // and notifies any connected servers. +// +// Deprecated: the roots feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). Migrate to passing paths via tool parameters, +// resource URIs, or configuration. See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. func (c *Client) AddRoots(roots ...*Root) { // Only notify if something could change. if len(roots) == 0 { @@ -256,6 +675,12 @@ func (c *Client) AddRoots(roots ...*Root) { // RemoveRoots removes the roots with the given URIs, // and notifies any connected servers if the list has changed. // It is not an error to remove a nonexistent root. +// +// Deprecated: the roots feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). Migrate to passing paths via tool parameters, +// resource URIs, or configuration. See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. func (c *Client) RemoveRoots(uris ...string) { changeAndNotify(c, notificationRootsListChanged, &RootsListChangedParams{}, func() bool { return c.roots.remove(uris...) }) @@ -269,10 +694,36 @@ func changeAndNotify[P Params](c *Client, notification string, params P, change // Lock for the change, but not for the notification. c.mu.Lock() if change() { - sessions = slices.Clone(c.sessions) + // Check if listChanged is enabled for this notification type. + if c.shouldSendListChangedNotification(notification) { + sessions = slices.Clone(c.sessions) + } } c.mu.Unlock() - notifySessions(sessions, notification, params) + notifySessions(sessions, notification, params, c.opts.Logger) +} + +// shouldSendListChangedNotification checks if the client's capabilities allow +// sending the given list-changed notification. +func (c *Client) shouldSendListChangedNotification(notification string) bool { + // Get effective capabilities (considering user-provided defaults). + caps := c.opts.Capabilities + + switch notification { + case notificationRootsListChanged: + // If user didn't specify capabilities, default behavior sends notifications. + if caps == nil { + return true + } + // Check RootsV2 first (preferred), then fall back to Roots. + if caps.RootsV2 != nil { + return caps.RootsV2.ListChanged + } + return caps.Roots.ListChanged + default: + // Unknown notification, allow by default. + return true + } } func (c *Client) listRoots(_ context.Context, req *ListRootsRequest) (*ListRootsResult, error) { @@ -287,47 +738,184 @@ func (c *Client) listRoots(_ context.Context, req *ListRootsRequest) (*ListRoots }, nil } -func (c *Client) createMessage(ctx context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) { - if c.opts.CreateMessageHandler == nil { - // TODO: wrap or annotate this error? Pick a standard code? - return nil, jsonrpc2.NewError(codeUnsupportedMethod, "client does not support CreateMessage") +func (c *Client) createMessage(ctx context.Context, req *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + if c.opts.CreateMessageWithToolsHandler != nil { + return c.opts.CreateMessageWithToolsHandler(ctx, req) + } + if c.opts.CreateMessageHandler != nil { + // Downconvert the request for the basic handler. + baseParams, err := req.Params.toBase() + if err != nil { + return nil, err + } + baseReq := &CreateMessageRequest{ + Session: req.Session, + Params: baseParams, + } + res, err := c.opts.CreateMessageHandler(ctx, baseReq) + if err != nil { + return nil, err + } + return res.toWithTools(), nil + } + return nil, &jsonrpc.Error{Code: codeUnsupportedMethod, Message: "client does not support CreateMessage"} +} + +// urlElicitationMiddleware returns middleware that automatically handles URL elicitation +// required errors by executing the elicitation handler, waiting for completion notifications, +// and retrying the operation. +// +// This middleware should be added to clients that want automatic URL elicitation handling: +// +// client := mcp.NewClient(impl, opts) +// client.AddSendingMiddleware(mcp.urlElicitationMiddleware()) +// +// TODO(rfindley): this isn't strictly necessary for the SEP, but may be +// useful. Propose exporting it. +func urlElicitationMiddleware() Middleware { + return func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + // Call the underlying handler. + res, err := next(ctx, method, req) + if err == nil { + return res, nil + } + + // Check if this is a URL elicitation required error. + var rpcErr *jsonrpc.Error + if !errors.As(err, &rpcErr) || rpcErr.Code != CodeURLElicitationRequired { + return res, err + } + + // Notifications don't support retries. + if strings.HasPrefix(method, "notifications/") { + return res, err + } + + // Extract the client session. + cs, ok := req.GetSession().(*ClientSession) + if !ok { + return res, err + } + + // Check if the client has an elicitation handler. + if cs.client.opts.ElicitationHandler == nil { + return res, err + } + + // Parse the elicitations from the error data. + var errorData struct { + Elicitations []*ElicitParams `json:"elicitations"` + } + if rpcErr.Data != nil { + if err := json.Unmarshal(rpcErr.Data, &errorData); err != nil { + return nil, fmt.Errorf("failed to parse URL elicitation error data: %w", err) + } + } + + // Validate that all elicitations are URL mode. + for _, elicit := range errorData.Elicitations { + mode := elicit.Mode + if mode == "" { + mode = "form" // Default mode. + } + if mode != "url" { + return nil, fmt.Errorf("URLElicitationRequired error must only contain URL mode elicitations, got %q", mode) + } + } + + // Register waiters for all elicitations before executing handlers + // to avoid race condition where notification arrives before waiter is registered. + type waiter struct { + await func(context.Context) error + cleanup func() + } + waiters := make([]waiter, 0, len(errorData.Elicitations)) + for _, elicitParams := range errorData.Elicitations { + await, cleanup := cs.registerElicitationWaiter(elicitParams.ElicitationID) + waiters = append(waiters, waiter{await: await, cleanup: cleanup}) + } + + // Ensure cleanup happens even if we return early. + defer func() { + for _, w := range waiters { + w.cleanup() + } + }() + + // Execute the elicitation handler for each elicitation. + for _, elicitParams := range errorData.Elicitations { + elicitReq := newClientRequest(cs, elicitParams) + _, elicitErr := cs.client.elicit(ctx, elicitReq) + if elicitErr != nil { + return nil, fmt.Errorf("URL elicitation failed: %w", elicitErr) + } + } + + // Wait for all elicitations to complete. + for _, w := range waiters { + if err := w.await(ctx); err != nil { + return nil, err + } + } + + // All elicitations complete, retry the original operation. + return next(ctx, method, req) + } } - return c.opts.CreateMessageHandler(ctx, req) } func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult, error) { if c.opts.ElicitationHandler == nil { - // TODO: wrap or annotate this error? Pick a standard code? - return nil, jsonrpc2.NewError(codeUnsupportedMethod, "client does not support elicitation") + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: "client does not support elicitation"} } - // Validate that the requested schema only contains top-level properties without nesting - schema, err := validateElicitSchema(req.Params.RequestedSchema) - if err != nil { - return nil, jsonrpc2.NewError(codeInvalidParams, err.Error()) + // Validate the elicitation parameters based on the mode. + mode := req.Params.Mode + if mode == "" { + mode = "form" } - res, err := c.opts.ElicitationHandler(ctx, req) - if err != nil { - return nil, err - } - - // Validate elicitation result content against requested schema - if schema != nil && res.Content != nil { - // TODO: is this the correct behavior if validation fails? - // It isn't the *server's* params that are invalid, so why would we return - // this code to the server? - resolved, err := schema.Resolve(nil) + switch mode { + case "form": + if req.Params.URL != "" { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: "URL must not be set for form elicitation"} + } + schema, err := validateElicitSchema(req.Params.RequestedSchema) if err != nil { - return nil, jsonrpc2.NewError(codeInvalidParams, fmt.Sprintf("failed to resolve requested schema: %v", err)) + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: err.Error()} } - - if err := resolved.Validate(res.Content); err != nil { - return nil, jsonrpc2.NewError(codeInvalidParams, fmt.Sprintf("elicitation result content does not match requested schema: %v", err)) + res, err := c.opts.ElicitationHandler(ctx, req) + if err != nil { + return nil, err } + // Validate elicitation result content against requested schema. + if res.Action == "accept" && schema != nil && res.Content != nil { + resolved, err := schema.Resolve(nil) + if err != nil { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("failed to resolve requested schema: %v", err)} + } + if err := resolved.Validate(res.Content); err != nil { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("elicitation result content does not match requested schema: %v", err)} + } + err = resolved.ApplyDefaults(&res.Content) + if err != nil { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("failed to apply schema defaults to elicitation result: %v", err)} + } + } + return res, nil + case "url": + if req.Params.RequestedSchema != nil { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: "requestedSchema must not be set for URL elicitation"} + } + if req.Params.URL == "" { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: "URL must be set for URL elicitation"} + } + // No schema validation for URL mode, just pass through to handler. + return c.opts.ElicitationHandler(ctx, req) + default: + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("unsupported elicitation mode: %q", mode)} } - - return res, nil } // validateElicitSchema validates that the schema conforms to MCP elicitation schema requirements. @@ -341,6 +929,9 @@ func validateElicitSchema(wireSchema any) (*jsonschema.Schema, error) { if err := remarshal(wireSchema, &schema); err != nil { return nil, err } + if schema == nil { + return nil, nil + } // The root schema must be of type "object" if specified if schema.Type != "" && schema.Type != "object" { @@ -369,7 +960,6 @@ func validateElicitProperty(propName string, propSchema *jsonschema.Schema) erro if len(propSchema.Properties) > 0 { return fmt.Errorf("elicit schema property %q contains nested properties, only primitive properties are allowed", propName) } - // Validate based on the property type - only primitives are supported switch propSchema.Type { case "string": @@ -378,8 +968,10 @@ func validateElicitProperty(propName string, propSchema *jsonschema.Schema) erro return validateElicitNumberProperty(propName, propSchema) case "boolean": return validateElicitBooleanProperty(propName, propSchema) + case "array": + return validateElicitArrayProperty(propName, propSchema) default: - return fmt.Errorf("elicit schema property %q has unsupported type %q, only string, number, integer, and boolean are allowed", propName, propSchema.Type) + return fmt.Errorf("elicit schema property %q has unsupported type %q, only string, number, integer, boolean, and array are allowed", propName, propSchema.Type) } } @@ -392,7 +984,7 @@ func validateElicitStringProperty(propName string, propSchema *jsonschema.Schema return fmt.Errorf("elicit schema property %q has enum values but type is %q, enums are only supported for string type", propName, propSchema.Type) } // Enum values themselves are validated by the JSON schema library - // Validate enumNames if present - must match enum length + // Validate legacy enumNames if present - must match enum length. if propSchema.Extra != nil { if enumNamesRaw, exists := propSchema.Extra["enumNames"]; exists { // Type check enumNames - should be a slice @@ -407,6 +999,15 @@ func validateElicitStringProperty(propName string, propSchema *jsonschema.Schema } return nil } + // Handle new style of titled enums. + if propSchema.OneOf != nil { + for _, entry := range propSchema.OneOf { + if err := validateTitledEnumEntry(entry); err != nil { + return fmt.Errorf("elicit schema property %q oneOf has invalid entry: %v", propName, err) + } + } + return nil + } // Validate format if specified - only specific formats are allowed if propSchema.Format != "" { @@ -439,7 +1040,7 @@ func validateElicitStringProperty(propName string, propSchema *jsonschema.Schema } } - return nil + return validateDefaultProperty[string](propName, propSchema) } // validateElicitNumberProperty validates number and integer-type properties. @@ -450,19 +1051,75 @@ func validateElicitNumberProperty(propName string, propSchema *jsonschema.Schema } } + intDefaultError := validateDefaultProperty[int](propName, propSchema) + floatDefaultError := validateDefaultProperty[float64](propName, propSchema) + if intDefaultError != nil && floatDefaultError != nil { + return fmt.Errorf("elicit schema property %q has default value that cannot be interpreted as an int or float", propName) + } + + return nil +} + +// validateElicitArrayProperty validates multi-select enum properties. +func validateElicitArrayProperty(propName string, propSchema *jsonschema.Schema) error { + if propSchema.Items == nil { + return fmt.Errorf("elicit schema property %q is array but missing 'items' definition", propName) + } + + items := propSchema.Items + switch items.Type { + case "string": + // Untitled enums. + if items.Enum == nil { + return fmt.Errorf("elicit schema property %q items must specify enum for untitled enums", propName) + } + return nil + case "": + // Titled enums. + if len(items.AnyOf) == 0 { + return fmt.Errorf("elicit schema property %q items must specify anyOf for titled enums", propName) + } + for _, entry := range items.AnyOf { + if err := validateTitledEnumEntry(entry); err != nil { + return fmt.Errorf("elicit schema property %q items has invalid entry: %v", propName, err) + } + } + return nil + default: + return fmt.Errorf("elicit schema property %q items have unsupported type %q", propName, items.Type) + } +} + +func validateTitledEnumEntry(entry *jsonschema.Schema) error { + if entry.Const == nil { + return fmt.Errorf("const is required for titled enum entries") + } + constVal, ok := (*entry.Const).(string) + if !ok { + return fmt.Errorf("const must be a string for titled enum entries") + } + if constVal == "" { + return fmt.Errorf("const cannot be empty for titled enum entries") + } + if entry.Title == "" { + return fmt.Errorf("title is required for titled enum entries") + } return nil } // validateElicitBooleanProperty validates boolean-type properties. func validateElicitBooleanProperty(propName string, propSchema *jsonschema.Schema) error { - // Validate default value if specified - must be a valid boolean + return validateDefaultProperty[bool](propName, propSchema) +} + +func validateDefaultProperty[T any](propName string, propSchema *jsonschema.Schema) error { + // Validate default value if specified - must be a valid T if propSchema.Default != nil { - var defaultValue bool + var defaultValue T if err := json.Unmarshal(propSchema.Default, &defaultValue); err != nil { - return fmt.Errorf("elicit schema property %q has invalid default value, must be a boolean: %v", propName, err) + return fmt.Errorf("elicit schema property %q has invalid default value, must be a %T: %v", propName, defaultValue, err) } } - return nil } @@ -490,6 +1147,9 @@ func (c *Client) AddSendingMiddleware(middleware ...Middleware) { // // Receiving middleware is called when a request is received. It is useful for tasks // such as authentication, request logging and metrics. +// +// A received message need not carry params: use [HasParams] before inspecting +// [Request.GetParams]. func (c *Client) AddReceivingMiddleware(middleware ...Middleware) { c.mu.Lock() defer c.mu.Unlock() @@ -514,10 +1174,14 @@ var clientMethodInfos = map[string]methodInfo{ notificationResourceUpdated: newClientMethodInfo(clientMethod((*Client).callResourceUpdatedHandler), notification|missingParamsOK), notificationLoggingMessage: newClientMethodInfo(clientMethod((*Client).callLoggingHandler), notification), notificationProgress: newClientMethodInfo(clientSessionMethod((*ClientSession).callProgressNotificationHandler), notification), + notificationElicitationComplete: newClientMethodInfo(clientMethod((*Client).callElicitationCompleteHandler), notification|missingParamsOK), + notificationSubscriptionsAck: newClientMethodInfo(clientMethod((*Client).callSubscriptionsAckHandler), notification|missingParamsOK), } func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo { - return serverMethodInfos + cs.client.mu.Lock() + defer cs.client.mu.Unlock() + return cs.client.sendMethods } func (cs *ClientSession) receivingMethodInfos() map[string]methodInfo { @@ -570,18 +1234,51 @@ func (cs *ClientSession) Ping(ctx context.Context, params *PingParams) error { } // ListPrompts lists prompts that are currently available on the server. +// +// Results may be served from a client-side TTL cache populated by previous +// calls; see SEP-2549. func (cs *ClientSession) ListPrompts(ctx context.Context, params *ListPromptsParams) (*ListPromptsResult, error) { - return handleSend[*ListPromptsResult](ctx, methodListPrompts, newClientRequest(cs, orZero[Params](params))) + if cs.usesNewProtocol() { + if result, ok := cachedListResult(&cs.promptsCache, params); ok { + return result, nil + } + params = injectRequestMeta(cs, params) + } + result, err := handleSend[*ListPromptsResult](ctx, methodListPrompts, newClientRequest(cs, orZero[Params](params))) + if err != nil { + return nil, err + } + if cs.usesNewProtocol() { + cs.promptsCache.put(params.Cursor, result) + } + return result, nil } // GetPrompt gets a prompt from the server. func (cs *ClientSession) GetPrompt(ctx context.Context, params *GetPromptParams) (*GetPromptResult, error) { + if cs.usesNewProtocol() { + params = injectRequestMeta(cs, params) + } return handleSend[*GetPromptResult](ctx, methodGetPrompt, newClientRequest(cs, orZero[Params](params))) } // ListTools lists tools that are currently available on the server. func (cs *ClientSession) ListTools(ctx context.Context, params *ListToolsParams) (*ListToolsResult, error) { - return handleSend[*ListToolsResult](ctx, methodListTools, newClientRequest(cs, orZero[Params](params))) + if cs.usesNewProtocol() { + if result, ok := cachedListResult(&cs.toolsCache, params); ok { + return result, nil + } + params = injectRequestMeta(cs, params) + } + result, err := handleSend[*ListToolsResult](ctx, methodListTools, newClientRequest(cs, orZero[Params](params))) + if err != nil { + return nil, err + } + result.Tools = filterValidTools(cs.client.opts.Logger, result.Tools) + if cs.usesNewProtocol() { + cs.toolsCache.put(params.Cursor, result) + } + return result, nil } // CallTool calls the tool with the given parameters. @@ -595,9 +1292,23 @@ func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams) ( // Avoid sending nil over the wire. params.Arguments = map[string]any{} } + if tool := cs.lookupTool(params.Name); tool != nil { + ctx = context.WithValue(ctx, toolContextKey, tool) + } + if cs.usesNewProtocol() { + params = injectRequestMeta(cs, params) + } return handleSend[*CallToolResult](ctx, methodCallTool, newClientRequest(cs, orZero[Params](params))) } +// SetLoggingLevel sets the minimum severity level for log messages sent by +// the server. +// +// Deprecated: the logging feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). Migrate to consuming stderr output (for STDIO +// servers) or OpenTelemetry. See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. func (cs *ClientSession) SetLoggingLevel(ctx context.Context, params *SetLoggingLevelParams) error { _, err := handleSend[*emptyResult](ctx, methodSetLevel, newClientRequest(cs, orZero[Params](params))) return err @@ -605,38 +1316,161 @@ func (cs *ClientSession) SetLoggingLevel(ctx context.Context, params *SetLogging // ListResources lists the resources that are currently available on the server. func (cs *ClientSession) ListResources(ctx context.Context, params *ListResourcesParams) (*ListResourcesResult, error) { - return handleSend[*ListResourcesResult](ctx, methodListResources, newClientRequest(cs, orZero[Params](params))) + if cs.usesNewProtocol() { + if result, ok := cachedListResult(&cs.resourcesCache, params); ok { + return result, nil + } + params = injectRequestMeta(cs, params) + } + result, err := handleSend[*ListResourcesResult](ctx, methodListResources, newClientRequest(cs, orZero[Params](params))) + if err != nil { + return nil, err + } + if cs.usesNewProtocol() { + cs.resourcesCache.put(params.Cursor, result) + } + return result, nil } // ListResourceTemplates lists the resource templates that are currently available on the server. func (cs *ClientSession) ListResourceTemplates(ctx context.Context, params *ListResourceTemplatesParams) (*ListResourceTemplatesResult, error) { - return handleSend[*ListResourceTemplatesResult](ctx, methodListResourceTemplates, newClientRequest(cs, orZero[Params](params))) + if cs.usesNewProtocol() { + if result, ok := cachedListResult(&cs.resourceTemplatesCache, params); ok { + return result, nil + } + params = injectRequestMeta(cs, params) + } + result, err := handleSend[*ListResourceTemplatesResult](ctx, methodListResourceTemplates, newClientRequest(cs, orZero[Params](params))) + if err != nil { + return nil, err + } + if cs.usesNewProtocol() { + cs.resourceTemplatesCache.put(params.Cursor, result) + } + return result, nil } // ReadResource asks the server to read a resource and return its contents. func (cs *ClientSession) ReadResource(ctx context.Context, params *ReadResourceParams) (*ReadResourceResult, error) { - return handleSend[*ReadResourceResult](ctx, methodReadResource, newClientRequest(cs, orZero[Params](params))) + if cs.usesNewProtocol() { + var uri string + if params != nil { + uri = params.URI + } + if result, ok := cs.readResourceCache.get(uri); ok { + return result, nil + } + params = injectRequestMeta(cs, params) + } + result, err := handleSend[*ReadResourceResult](ctx, methodReadResource, newClientRequest(cs, orZero[Params](params))) + if err != nil { + return nil, err + } + if cs.usesNewProtocol() { + cs.readResourceCache.put(params.URI, result) + } + return result, nil } func (cs *ClientSession) Complete(ctx context.Context, params *CompleteParams) (*CompleteResult, error) { + if cs.usesNewProtocol() { + params = injectRequestMeta(cs, params) + } return handleSend[*CompleteResult](ctx, methodComplete, newClientRequest(cs, orZero[Params](params))) } // Subscribe sends a "resources/subscribe" request to the server, asking for // notifications when the specified resource changes. func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams) error { - _, err := handleSend[*emptyResult](ctx, methodSubscribe, newClientRequest(cs, orZero[Params](params))) + if !cs.usesNewProtocol() { + _, err := handleSend[*emptyResult](ctx, methodSubscribe, newClientRequest(cs, orZero[Params](params))) + return err + } + if params == nil || params.URI == "" { + return fmt.Errorf("Subscribe: missing URI") + } + uri := params.URI + + var listenCtx context.Context + cs.resourceSubsMu.Lock() + if _, exists := cs.resourceSubs[uri]; !exists { + var cancel context.CancelFunc + listenCtx, cancel = context.WithCancel(context.Background()) + if cs.resourceSubs == nil { + cs.resourceSubs = make(map[string]context.CancelFunc) + } + cs.resourceSubs[uri] = cancel + } + cs.resourceSubsMu.Unlock() + if listenCtx == nil { + // Already subscribed to this URI + return nil + } + + return cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{ + Notifications: &NotificationSubscriptions{ + ResourceSubscriptions: []string{uri}, + }, + }) +} + +// Unsubscribe cancels a previous [ClientSession.Subscribe] for params.URI. +// +// Under the legacy protocol it sends a "resources/unsubscribe" request. +// +// Under SEP-2575 it cancels the background "subscriptions/listen" stream +// opened by Subscribe for the URI. Unsubscribe is idempotent: calling it for +// a URI that is not currently subscribed is a no-op. +func (cs *ClientSession) Unsubscribe(ctx context.Context, params *UnsubscribeParams) error { + if !cs.usesNewProtocol() { + _, err := handleSend[*emptyResult](ctx, methodUnsubscribe, newClientRequest(cs, orZero[Params](params))) + return err + } + if params == nil || params.URI == "" { + return fmt.Errorf("Unsubscribe: missing URI") + } + cs.resourceSubsMu.Lock() + cancel, ok := cs.resourceSubs[params.URI] + delete(cs.resourceSubs, params.URI) + cs.resourceSubsMu.Unlock() + if ok { + cancel() + } + return nil +} + +// cancelAllResourceSubscriptions cancels every active SEP-2575 resource +// subscription opened via Subscribe. The listen goroutines exit +// asynchronously as their contexts unwind. Called from Close. +func (cs *ClientSession) cancelAllResourceSubscriptions() { + cs.resourceSubsMu.Lock() + subs := cs.resourceSubs + cs.resourceSubs = nil + cs.resourceSubsMu.Unlock() + for _, cancel := range subs { + cancel() + } +} + +// SubscriptionsListen opens a SEP-2575 "subscriptions/listen" stream. +// +// The server's first message on the stream is "notifications/subscriptions/acknowledged"; +// subsequent opted-in notifications (e.g. tools/list_changed) are delivered through the +// usual handlers registered in [ClientOptions]. +func (cs *ClientSession) subscriptionsListen(ctx context.Context, params *SubscriptionsListenParams) error { + params = injectRequestMeta(cs, params) + _, err := handleSend[*SubscriptionsListenResult](ctx, methodSubscriptionsListen, newClientRequest(cs, orZero[Params](params))) return err } -// Unsubscribe sends a "resources/unsubscribe" request to the server, cancelling -// a previous subscription. -func (cs *ClientSession) Unsubscribe(ctx context.Context, params *UnsubscribeParams) error { - _, err := handleSend[*emptyResult](ctx, methodUnsubscribe, newClientRequest(cs, orZero[Params](params))) - return err +func (c *Client) callSubscriptionsAckHandler(context.Context, *ClientRequest[*SubscriptionsAcknowledgedParams]) (Result, error) { + return nil, nil } func (c *Client) callToolChangedHandler(ctx context.Context, req *ToolListChangedRequest) (Result, error) { + if cs, ok := req.GetSession().(*ClientSession); ok { + cs.toolsCache.invalidate() + } if h := c.opts.ToolListChangedHandler; h != nil { h(ctx, req) } @@ -644,6 +1478,9 @@ func (c *Client) callToolChangedHandler(ctx context.Context, req *ToolListChange } func (c *Client) callPromptChangedHandler(ctx context.Context, req *PromptListChangedRequest) (Result, error) { + if cs, ok := req.GetSession().(*ClientSession); ok { + cs.promptsCache.invalidate() + } if h := c.opts.PromptListChangedHandler; h != nil { h(ctx, req) } @@ -651,6 +1488,10 @@ func (c *Client) callPromptChangedHandler(ctx context.Context, req *PromptListCh } func (c *Client) callResourceChangedHandler(ctx context.Context, req *ResourceListChangedRequest) (Result, error) { + if cs, ok := req.GetSession().(*ClientSession); ok { + cs.resourcesCache.invalidate() + cs.resourceTemplatesCache.invalidate() + } if h := c.opts.ResourceListChangedHandler; h != nil { h(ctx, req) } @@ -658,6 +1499,9 @@ func (c *Client) callResourceChangedHandler(ctx context.Context, req *ResourceLi } func (c *Client) callResourceUpdatedHandler(ctx context.Context, req *ResourceUpdatedNotificationRequest) (Result, error) { + if cs, ok := req.GetSession().(*ClientSession); ok && req.Params != nil { + cs.readResourceCache.invalidateKey(req.Params.URI) + } if h := c.opts.ResourceUpdatedHandler; h != nil { h(ctx, req) } @@ -678,6 +1522,27 @@ func (cs *ClientSession) callProgressNotificationHandler(ctx context.Context, pa return nil, nil } +func (c *Client) callElicitationCompleteHandler(ctx context.Context, req *ElicitationCompleteNotificationRequest) (Result, error) { + // Check if there's a pending elicitation waiting for this notification. + if cs, ok := req.GetSession().(*ClientSession); ok { + cs.pendingElicitationsMu.Lock() + if ch, exists := cs.pendingElicitations[req.Params.ElicitationID]; exists { + select { + case ch <- struct{}{}: + default: + // Channel already signaled. + } + } + cs.pendingElicitationsMu.Unlock() + } + + // Call the user's handler if provided. + if h := c.opts.ElicitationCompleteHandler; h != nil { + h(ctx, req) + } + return nil, nil +} + // NotifyProgress sends a progress notification from the client to the server // associated with this session. // This can be used if the client is performing a long-running task that was @@ -760,3 +1625,67 @@ func paginate[P listParams, R listResult[T], T any](ctx context.Context, params } } } + +// AddSendingCustomMethod registers a custom JSON-RPC method +// that the client may send to the server. +// +// Registration is decoupled from invocation: extensions typically call this +// during setup, while the actual call site uses [CallCustomMethod]. +// +// if err := mcp.AddSendingCustomMethod[*SearchParams, *SearchResult](c, "acme/search"); err != nil { +// return err +// } +// // ... later, anywhere a *ClientSession is available: +// result, err := mcp.CallCustomMethod[*SearchParams, *SearchResult]( +// ctx, cs, "acme/search", &SearchParams{Query: "hello"}) +// +// AddSendingCustomMethod returns an error if method is the name of a standard +// MCP method. Registering the same method twice replaces the previous +// registration. +func AddSendingCustomMethod[P paramsPtr[PT], R Result, PT any]( + c *Client, + method string, +) error { + if _, ok := serverMethodInfos[method]; ok { + return fmt.Errorf("mcp: AddSendingCustomMethod: %q shadows a standard MCP method", method) + } + + mi := methodInfo{ + newResult: func() Result { + return reflect.New(reflect.TypeFor[R]().Elem()).Interface().(R) + }, + } + + c.mu.Lock() + defer c.mu.Unlock() + c.sendMethods[method] = mi + return nil +} + +// CallCustomMethod sends a custom (non-standard) JSON-RPC method to the +// server and decodes the response into R. +// +// The method must have been registered on the session's client via +// [AddSendingCustomMethod]. +func CallCustomMethod[P paramsPtr[PT], R Result, PT any]( + ctx context.Context, + cs *ClientSession, + method string, + params P, +) (R, error) { + c := cs.client + c.mu.Lock() + _, ok := c.sendMethods[method] + c.mu.Unlock() + if !ok { + var zero R + return zero, fmt.Errorf("mcp: CallCustomMethod: %q is not registered; call AddSendingCustomMethod first", method) + } + if cs.usesNewProtocol() { + params = injectRequestMeta(cs, params) + } + return handleSend[R](ctx, method, &ClientRequest[P]{ + Session: cs, + Params: params, + }) +} diff --git a/mcp/client_example_test.go b/mcp/client_example_test.go index cc7146c..98ac3b1 100644 --- a/mcp/client_example_test.go +++ b/mcp/client_example_test.go @@ -6,6 +6,7 @@ package mcp_test import ( "context" + "encoding/json" "fmt" "log" @@ -18,26 +19,27 @@ import ( func Example_roots() { ctx := context.Background() - // Create a client with a single root. + // Create a client with two roots. c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) - c.AddRoots(&mcp.Root{URI: "file://a"}) + c.AddRoots(&mcp.Root{URI: "file://a"}, &mcp.Root{URI: "file://b"}) - // Now create a server with a handler to receive notifications about roots. - rootsChanged := make(chan struct{}) - handleRootsChanged := func(ctx context.Context, req *mcp.RootsListChangedRequest) { - rootList, err := req.Session.ListRoots(ctx, nil) - if err != nil { - log.Fatal(err) + // Create a server with a tool that requests roots via the multi round-trip + // pattern (SEP-2322): server-to-client requests are no longer sent as + // standalone JSON-RPC calls on protocol version >= 2026-07-28. + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(s, &mcp.Tool{Name: "roots"}, func(_ context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{"roots": &mcp.ListRootsParams{}}, + }, nil, nil } + rootList := req.Params.InputResponses["roots"].(*mcp.ListRootsResult) var roots []string for _, root := range rootList.Roots { roots = append(roots, root.URI) } fmt.Println(roots) - close(rootsChanged) - } - s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, &mcp.ServerOptions{ - RootsListChangedHandler: handleRootsChanged, + return &mcp.CallToolResult{}, nil, nil }) // Connect the server and client... @@ -54,14 +56,71 @@ func Example_roots() { } defer clientSession.Close() - // ...and add a root. The server is notified about the change. - c.AddRoots(&mcp.Root{URI: "file://b"}) - <-rootsChanged + // ...and call the tool. The client's multi round-trip driver fulfils the + // embedded roots/list request and retries the call. + if _, err := clientSession.CallTool(ctx, &mcp.CallToolParams{Name: "roots"}); err != nil { + log.Fatal(err) + } // Output: [file://a file://b] } // !-roots +// !+rootslistchanged + +func Example_rootsListChanged() { + ctx := context.Background() + + changed := make(chan struct{}, 2) + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, &mcp.ServerOptions{ + RootsListChangedHandler: func(context.Context, *mcp.RootsListChangedRequest) { + changed <- struct{}{} + }, + }) + + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) + c.AddRoots(&mcp.Root{URI: "file:///project"}) + + t1, t2 := mcp.NewInMemoryTransports() + ss, err := s.Connect(ctx, t1, nil) + if err != nil { + log.Fatal(err) + } + defer ss.Close() + + // ListRoots is a server-initiated request, so this session negotiates a + // protocol version that still allows one. + cs, err := c.Connect(ctx, t2, &mcp.ClientSessionOptions{ProtocolVersion: "2025-11-25"}) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + // Roots added after the client connects notify every connected server. + c.AddRoots(&mcp.Root{URI: "file:///scratch"}) + <-changed + + // The notification says only that the list changed, so read it back. + res, err := ss.ListRoots(ctx, nil) + if err != nil { + log.Fatal(err) + } + for _, root := range res.Roots { + fmt.Println(root.URI) + } + + c.RemoveRoots("file:///scratch") + <-changed + fmt.Println("roots changed again") + + // Output: + // file:///project + // file:///scratch + // roots changed again +} + +// !-rootslistchanged + // !+sampling func Example_sampling() { @@ -80,22 +139,35 @@ func Example_sampling() { // Connect the server and client... ct, st := mcp.NewInMemoryTransports() + // Create a server with a tool that requests sampling via the multi + // round-trip pattern (SEP-2322): server-to-client requests are no longer + // sent as standalone JSON-RPC calls on protocol version >= 2026-07-28. s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(s, &mcp.Tool{Name: "sample"}, func(_ context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{"msg": &mcp.CreateMessageParams{}}, + }, nil, nil + } + msg := req.Params.InputResponses["msg"].(*mcp.CreateMessageWithToolsResult) + return &mcp.CallToolResult{Content: msg.Content}, nil, nil + }) session, err := s.Connect(ctx, st, nil) if err != nil { log.Fatal(err) } defer session.Close() - if _, err := c.Connect(ctx, ct, nil); err != nil { - log.Fatal(err) - } - - msg, err := session.CreateMessage(ctx, &mcp.CreateMessageParams{}) + clientSession, err := c.Connect(ctx, ct, nil) if err != nil { log.Fatal(err) } - fmt.Println(msg.Content.(*mcp.TextContent).Text) + + res, err := clientSession.CallTool(ctx, &mcp.CallToolParams{Name: "sample"}) + if err != nil { + log.Fatal(err) + } + fmt.Println(res.Content[0].(*mcp.TextContent).Text) // Output: would have created a message } @@ -107,7 +179,28 @@ func Example_elicitation() { ctx := context.Background() ct, st := mcp.NewInMemoryTransports() + // Create a server with a tool that requests elicitation via the multi + // round-trip pattern (SEP-2322): server-to-client requests are no longer + // sent as standalone JSON-RPC calls on protocol version >= 2026-07-28. s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(s, &mcp.Tool{Name: "ask"}, func(_ context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{"input": &mcp.ElicitParams{ + Message: "This should fail", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "test": {Type: "string"}, + }, + }, + }}, + }, nil, nil + } + res := req.Params.InputResponses["input"].(*mcp.ElicitResult) + fmt.Println(res.Content["test"]) + return &mcp.CallToolResult{}, nil, nil + }) ss, err := s.Connect(ctx, st, nil) if err != nil { log.Fatal(err) @@ -119,23 +212,128 @@ func Example_elicitation() { return &mcp.ElicitResult{Action: "accept", Content: map[string]any{"test": "value"}}, nil }, }) - if _, err := c.Connect(ctx, ct, nil); err != nil { - log.Fatal(err) - } - res, err := ss.Elicit(ctx, &mcp.ElicitParams{ - Message: "This should fail", - RequestedSchema: &jsonschema.Schema{ - Type: "object", - Properties: map[string]*jsonschema.Schema{ - "test": {Type: "string"}, - }, - }, - }) + clientSession, err := c.Connect(ctx, ct, nil) if err != nil { log.Fatal(err) } - fmt.Println(res.Content["test"]) + if _, err := clientSession.CallTool(ctx, &mcp.CallToolParams{Name: "ask"}); err != nil { + log.Fatal(err) + } // Output: value } // !-elicitation + +// !+elicitationschema + +func Example_elicitationSchema() { + ctx := context.Background() + ct, st := mcp.NewInMemoryTransports() + + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(s, &mcp.Tool{Name: "export_report"}, func(_ context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{"format": &mcp.ElicitParams{ + Message: "Export quarterly-sales as which format?", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "format": { + Type: "string", + Title: "Format", + Enum: []any{"pdf", "csv"}, + Default: json.RawMessage(`"pdf"`), + Extra: map[string]any{"enumNames": []any{"PDF document", "CSV spreadsheet"}}, + }, + }, + }, + }}, + }, nil, nil + } + res := req.Params.InputResponses["format"].(*mcp.ElicitResult) + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "Exported as " + res.Content["format"].(string)}}, + }, nil, nil + }) + if _, err := s.Connect(ctx, st, nil); err != nil { + log.Fatal(err) + } + + // The user accepts without filling anything in. + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, &mcp.ClientOptions{ + ElicitationHandler: func(context.Context, *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + return &mcp.ElicitResult{Action: "accept", Content: map[string]any{}}, nil + }, + }) + cs, err := c.Connect(ctx, ct, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + res, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: "export_report"}) + if err != nil { + log.Fatal(err) + } + fmt.Println(res.Content[0].(*mcp.TextContent).Text) + // Output: Exported as pdf +} + +// !-elicitationschema + +// !+elicitationcomplete + +func Example_elicitationComplete() { + ctx := context.Background() + ct, st := mcp.NewInMemoryTransports() + + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + log.Fatal(err) + } + defer ss.Close() + + done := make(chan struct{}) + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, &mcp.ClientOptions{ + Capabilities: &mcp.ClientCapabilities{ + Elicitation: &mcp.ElicitationCapabilities{URL: &mcp.URLElicitationCapabilities{}}, + }, + ElicitationHandler: func(_ context.Context, req *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + fmt.Println("opening", req.Params.URL) + return &mcp.ElicitResult{Action: "accept"}, nil + }, + ElicitationCompleteHandler: func(_ context.Context, req *mcp.ElicitationCompleteNotificationRequest) { + fmt.Println("flow finished:", req.Params.ElicitationID) + close(done) + }, + }) + cs, err := c.Connect(ctx, ct, &mcp.ClientSessionOptions{ProtocolVersion: "2025-11-25"}) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + const elicitationID = "connect-calendar-1" + if _, err := ss.Elicit(ctx, &mcp.ElicitParams{ + Message: "Grant calendar access", + URL: "https://calendar.example.com/consent?state=" + elicitationID, + ElicitationID: elicitationID, + }); err != nil { + log.Fatal(err) + } + + // The hosted page redirects back to the server, whose callback endpoint + // signals that the user is done. + if err := ss.NotifyElicitationComplete(ctx, &mcp.ElicitationCompleteParams{ElicitationID: elicitationID}); err != nil { + log.Fatal(err) + } + <-done + + // Output: + // opening https://calendar.example.com/consent?state=connect-calendar-1 + // flow finished: connect-calendar-1 +} + +// !-elicitationcomplete diff --git a/mcp/client_test.go b/mcp/client_test.go index eaeedc8..f02b905 100644 --- a/mcp/client_test.go +++ b/mcp/client_test.go @@ -2,16 +2,22 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. +//lint:file-ignore SA1019 tests exercise deprecated SEP-2577 APIs (roots, sampling, logging). + package mcp import ( "context" "fmt" + "log/slog" + "sync/atomic" "testing" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) type Item struct { @@ -165,6 +171,23 @@ func TestClientPaginateBasic(t *testing.T) { } } +func TestClientLogger(t *testing.T) { + // Case 1: No logger provided + c1 := NewClient(&Implementation{Name: "test", Version: "1.0"}, nil) + if c1.opts.Logger == nil { + t.Error("expected default logger, got nil") + } + + // Case 2: Logger provided + logger := slog.Default() + c2 := NewClient(&Implementation{Name: "test", Version: "1.0"}, &ClientOptions{ + Logger: logger, + }) + if c2.opts.Logger != logger { + t.Error("expected provided logger, got different one") + } +} + func TestClientPaginateVariousPageSizes(t *testing.T) { ctx := context.Background() for i := 1; i < len(allItems)+1; i++ { @@ -196,19 +219,19 @@ func TestClientCapabilities(t *testing.T) { name string configureClient func(s *Client) clientOpts ClientOptions + protocolVersion string // defaults to latestProtocolVersion if empty wantCapabilities *ClientCapabilities }{ { - name: "With initial capabilities", + name: "default", configureClient: func(s *Client) {}, wantCapabilities: &ClientCapabilities{ - Roots: struct { - ListChanged bool "json:\"listChanged,omitempty\"" - }{ListChanged: true}, + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, }, }, { - name: "With sampling", + name: "with sampling", configureClient: func(s *Client) {}, clientOpts: ClientOptions{ CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { @@ -216,22 +239,856 @@ func TestClientCapabilities(t *testing.T) { }, }, wantCapabilities: &ClientCapabilities{ - Roots: struct { - ListChanged bool "json:\"listChanged,omitempty\"" - }{ListChanged: true}, + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, Sampling: &SamplingCapabilities{}, }, }, + { + name: "with elicitation", + configureClient: func(s *Client) {}, + clientOpts: ClientOptions{ + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return nil, nil + }, + }, + protocolVersion: protocolVersion20251125, + wantCapabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Elicitation: &ElicitationCapabilities{ + Form: &FormElicitationCapabilities{}, + }, + }, + }, + { + name: "with elicitation (old protocol)", + configureClient: func(s *Client) {}, + clientOpts: ClientOptions{ + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return nil, nil + }, + }, + protocolVersion: protocolVersion20250618, + wantCapabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Elicitation: &ElicitationCapabilities{}, + }, + }, + { + name: "with URL elicitation", + configureClient: func(s *Client) {}, + clientOpts: ClientOptions{ + Capabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Elicitation: &ElicitationCapabilities{ + URL: &URLElicitationCapabilities{}, + }, + }, + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return nil, nil + }, + }, + wantCapabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Elicitation: &ElicitationCapabilities{ + URL: &URLElicitationCapabilities{}, + }, + }, + }, + { + name: "with form and URL elicitation", + configureClient: func(s *Client) {}, + clientOpts: ClientOptions{ + Capabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Elicitation: &ElicitationCapabilities{ + Form: &FormElicitationCapabilities{}, + URL: &URLElicitationCapabilities{}, + }, + }, + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return nil, nil + }, + }, + wantCapabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Elicitation: &ElicitationCapabilities{ + Form: &FormElicitationCapabilities{}, + URL: &URLElicitationCapabilities{}, + }, + }, + }, + { + name: "no capabilities", + configureClient: func(s *Client) {}, + clientOpts: ClientOptions{ + Capabilities: &ClientCapabilities{}, + }, + wantCapabilities: &ClientCapabilities{}, + }, + { + name: "no roots", + configureClient: func(s *Client) {}, + clientOpts: ClientOptions{ + Capabilities: &ClientCapabilities{ + Sampling: &SamplingCapabilities{}, + }, + }, + wantCapabilities: &ClientCapabilities{ + Sampling: &SamplingCapabilities{}, + }, + }, + { + name: "roots-no list", + configureClient: func(s *Client) {}, + clientOpts: ClientOptions{ + Capabilities: &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: false}, + }, + }, + wantCapabilities: &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: false}, + }, + }, + { + name: "custom capabilities with sampling", + configureClient: func(s *Client) {}, + clientOpts: ClientOptions{ + Capabilities: &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: true}, + }, + CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + return nil, nil + }, + }, + wantCapabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Sampling: &SamplingCapabilities{}, + }, + }, + { + name: "elicitation override", + configureClient: func(s *Client) {}, + clientOpts: ClientOptions{ + Capabilities: &ClientCapabilities{ + Elicitation: &ElicitationCapabilities{ + URL: &URLElicitationCapabilities{}, + }, + }, + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return nil, nil + }, + }, + wantCapabilities: &ClientCapabilities{ + Elicitation: &ElicitationCapabilities{ + URL: &URLElicitationCapabilities{}, + }, + }, + }, + { + name: "custom capabilities with experimental", + configureClient: func(s *Client) {}, + clientOpts: ClientOptions{ + Capabilities: &ClientCapabilities{ + Experimental: map[string]any{"custom": "value"}, + RootsV2: &RootCapabilities{ListChanged: true}, + }, + }, + wantCapabilities: &ClientCapabilities{ + Experimental: map[string]any{"custom": "value"}, + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + }, + }, + { + name: "extensions preserved", + configureClient: func(s *Client) {}, + clientOpts: func() ClientOptions { + caps := &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: true}, + } + caps.AddExtension("io.example/ext1", map[string]any{"key": "value"}) + caps.AddExtension("io.example/ext2", nil) + return ClientOptions{Capabilities: caps} + }(), + wantCapabilities: &ClientCapabilities{ + Extensions: map[string]any{ + "io.example/ext1": map[string]any{"key": "value"}, + "io.example/ext2": map[string]any{}, + }, + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + }, + }, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { client := NewClient(testImpl, &tc.clientOpts) tc.configureClient(client) - gotCapabilities := client.capabilities() + protocolVersion := tc.protocolVersion + if protocolVersion == "" { + protocolVersion = latestProtocolVersion + } + gotCapabilities := client.capabilities(protocolVersion) if diff := cmp.Diff(tc.wantCapabilities, gotCapabilities); diff != "" { t.Errorf("capabilities() mismatch (-want +got):\n%s", diff) } }) } } + +func TestLookupTool(t *testing.T) { + tool1 := &Tool{Name: "tool1", Description: "first"} + tool2 := &Tool{Name: "tool2", Description: "second"} + tool1Updated := &Tool{Name: "tool1", Description: "updated"} + + // page represents a single cached ListToolsResult, keyed by cursor. + type page struct { + cursor string + tools []*Tool + } + + testCases := []struct { + name string + pages []page + lookup string + want *Tool + }{ + { + name: "empty cache", + lookup: "tool1", + want: nil, + }, + { + name: "single tool found", + pages: []page{{cursor: "", tools: []*Tool{tool1}}}, + lookup: "tool1", + want: tool1, + }, + { + name: "unknown tool", + pages: []page{{cursor: "", tools: []*Tool{tool1}}}, + lookup: "nonexistent", + want: nil, + }, + { + name: "multiple tools single page", + pages: []page{{cursor: "", tools: []*Tool{tool1, tool2}}}, + lookup: "tool2", + want: tool2, + }, + { + name: "tool found across paginated pages", + pages: []page{ + {cursor: "", tools: []*Tool{tool1}}, + {cursor: "page2", tools: []*Tool{tool2}}, + }, + lookup: "tool2", + want: tool2, + }, + { + name: "re-list same cursor overwrites entry", + pages: []page{ + {cursor: "", tools: []*Tool{tool1}}, + {cursor: "", tools: []*Tool{tool1Updated}}, + }, + lookup: "tool1", + want: tool1Updated, + }, + { + name: "empty page no-op", + pages: []page{{cursor: "", tools: []*Tool{}}}, + lookup: "tool1", + want: nil, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + cs := &ClientSession{} + for _, p := range tc.pages { + cs.toolsCache.put(p.cursor, &ListToolsResult{ + Cacheable: Cacheable{TTLMs: 60_000}, + Tools: p.tools, + }) + } + got := cs.lookupTool(tc.lookup) + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Errorf("lookupTool(%q) mismatch (-want +got):\n%s", tc.lookup, diff) + } + }) + } +} + +func TestClientCapabilitiesOverWire(t *testing.T) { + testCases := []struct { + name string + clientOpts *ClientOptions + wantCapabilities *ClientCapabilities + }{ + { + name: "Default capabilities", + clientOpts: nil, + wantCapabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + }, + }, + { + name: "Custom Capabilities with roots listChanged false", + clientOpts: &ClientOptions{ + Capabilities: &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: false}, + }, + }, + wantCapabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: false}, + RootsV2: &RootCapabilities{ListChanged: false}, + }, + }, + { + name: "Dynamic sampling capability", + clientOpts: &ClientOptions{ + Capabilities: &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: true}, + }, + CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + return nil, nil + }, + }, + wantCapabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Sampling: &SamplingCapabilities{}, + }, + }, + { + name: "Empty capabilities disables defaults", + clientOpts: &ClientOptions{ + Capabilities: &ClientCapabilities{}, + }, + wantCapabilities: &ClientCapabilities{}, + }, + { + name: "Extensions over wire", + clientOpts: func() *ClientOptions { + caps := &ClientCapabilities{ + RootsV2: &RootCapabilities{ListChanged: true}, + } + caps.AddExtension("io.example/ext", map[string]any{"key": "value"}) + return &ClientOptions{Capabilities: caps} + }(), + wantCapabilities: &ClientCapabilities{ + Extensions: map[string]any{ + "io.example/ext": map[string]any{"key": "value"}, + }, + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + + // Create client. + impl := &Implementation{Name: "testClient", Version: "v1.0.0"} + client := NewClient(impl, tc.clientOpts) + + // Connect client and server. + cTransport, sTransport := NewInMemoryTransports() + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + ss, err := server.Connect(ctx, sTransport, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + cs, err := client.Connect(ctx, cTransport, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Check that the server received the expected capabilities. + initParams := ss.InitializeParams() + if initParams == nil { + t.Fatal("InitializeParams is nil") + } + + if diff := cmp.Diff(tc.wantCapabilities, initParams.Capabilities); diff != "" { + t.Errorf("Capabilities mismatch (-want +got):\n%s", diff) + } + }) + } +} + +// TestClientConnectDiscover exercises the SEP-2575 server/discover probe that +// Client.Connect now sends before falling back to the legacy initialize +// handshake. +// +// Each subtest installs a server-side receiving middleware that intercepts the +// "server/discover" method and returns a canned response: a successful +// DiscoverResult, a "Method not found" error, an UnsupportedProtocolVersion +// error, an unrelated failure, or a DiscoverResult whose supportedVersions +// don't overlap with the SDK. The test then asserts the resulting session +// state and whether the legacy initialize handshake ran. +func TestClientConnectDiscover(t *testing.T) { + const otherVersionsOnly = "1999-01-01" + + tests := []struct { + name string + // discoverHandler decides how the server replies to server/discover. + // Returning (nil, nil) means "let the default stub handle it" (which + // returns ErrMethodNotFound). + discoverHandler func() (Result, error) + // wantInitialize is true if the legacy initialize handshake should + // have run (i.e. discover signaled "not supported"). + wantInitialize bool + // wantVersion is the protocol version expected to end up on + // ClientSession.state.InitializeResult after Connect returns. + wantVersion string + }{ + { + name: "discover success skips initialize", + discoverHandler: func() (Result, error) { + return &DiscoverResult{ + Meta: Meta{ + MetaKeyServerInfo: &Implementation{Name: "discoverServer", Version: "v1.0.0"}, + }, + SupportedVersions: []string{protocolVersion20260728}, + Capabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: true}, + }, + }, nil + }, + wantInitialize: false, + wantVersion: protocolVersion20260728, + }, + { + name: "method not found falls back to initialize", + discoverHandler: func() (Result, error) { + return nil, jsonrpc2.ErrMethodNotFound + }, + wantInitialize: true, + wantVersion: protocolVersion20251125, + }, + { + name: "unsupported protocol version falls back to initialize", + discoverHandler: func() (Result, error) { + return nil, &jsonrpc.Error{ + Code: CodeUnsupportedProtocolVersion, + Message: "unsupported protocol version", + } + }, + wantInitialize: true, + wantVersion: protocolVersion20251125, + }, + { + name: "no overlapping supported version falls back to initialize", + discoverHandler: func() (Result, error) { + return &DiscoverResult{ + Meta: Meta{ + MetaKeyServerInfo: &Implementation{Name: "discoverServer", Version: "v1.0.0"}, + }, + SupportedVersions: []string{otherVersionsOnly}, + Capabilities: &ServerCapabilities{}, + }, nil + }, + wantInitialize: true, + wantVersion: protocolVersion20251125, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + + var ( + gotDiscover atomic.Bool + gotInitialize atomic.Bool + ) + + s := NewServer(testImpl, nil) + s.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + switch method { + case methodDiscover: + gotDiscover.Store(true) + return tc.discoverHandler() + case methodInitialize: + gotInitialize.Store(true) + } + return next(ctx, method, req) + } + }) + + ct, st := NewInMemoryTransports() + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server Connect: %v", err) + } + defer ss.Close() + + c := NewClient(testImpl, nil) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer cs.Close() + + if !gotDiscover.Load() { + t.Error("server did not receive server/discover") + } + if got, want := gotInitialize.Load(), tc.wantInitialize; got != want { + t.Errorf("initialize invoked = %v, want %v", got, want) + } + ir := cs.InitializeResult() + if ir == nil { + t.Fatal("InitializeResult is nil after Connect") + } + if got, want := ir.ProtocolVersion, tc.wantVersion; got != want { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q", got, want) + } + }) + } +} + +// TestClientConnectDiscover_RequestContents verifies that the server/discover +// request sent by Client.Connect carries the SEP-2575 per-request _meta triple: +// protocolVersion, clientInfo, and clientCapabilities. +func TestClientConnectDiscover_RequestContents(t *testing.T) { + ctx := context.Background() + + type captured struct { + params *DiscoverParams + } + var got captured + + s := NewServer(testImpl, nil) + s.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == methodDiscover { + sr, ok := req.(*ServerRequest[*DiscoverParams]) + if !ok { + t.Errorf("discover req has unexpected type %T", req) + return nil, jsonrpc2.ErrMethodNotFound + } + got.params = sr.Params + // Make discover "not supported" so Connect proceeds (we only + // care about the discover request payload here). + return nil, jsonrpc2.ErrMethodNotFound + } + return next(ctx, method, req) + } + }) + + ct, st := NewInMemoryTransports() + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server Connect: %v", err) + } + defer ss.Close() + + clientImpl := &Implementation{Name: "discover-probe-client", Version: "v9.9.9"} + c := NewClient(clientImpl, &ClientOptions{ + CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + return nil, nil + }, + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer cs.Close() + + if got.params == nil { + t.Fatal("server did not receive server/discover") + } + + meta := got.params.GetMeta() + if v, _ := meta[MetaKeyProtocolVersion].(string); v != protocolVersion20260728 { + t.Errorf("_meta[%s] = %q, want %q", MetaKeyProtocolVersion, v, protocolVersion20260728) + } + // _meta values round-trip through JSON, so on the server side they + // arrive as map[string]any rather than the typed Go pointers we sent. + info, ok := meta[MetaKeyClientInfo].(map[string]any) + if !ok { + t.Fatalf("_meta[%s] = %T, want map[string]any", MetaKeyClientInfo, meta[MetaKeyClientInfo]) + } + if got, want := info["name"], any(clientImpl.Name); got != want { + t.Errorf("clientInfo.name = %v, want %v", got, want) + } + caps, ok := meta[MetaKeyClientCapabilities].(map[string]any) + if !ok { + t.Fatalf("_meta[%s] = %T, want map[string]any", MetaKeyClientCapabilities, meta[MetaKeyClientCapabilities]) + } + if _, ok := caps["sampling"]; !ok { + t.Errorf("clientCapabilities.sampling missing (CreateMessageHandler was set); got %v", caps) + } +} + +// TestInMemory_E2E_DiscoverSuccess is a full end-to-end smoke test for +// SEP-2575 over a non-HTTP transport. +func TestInMemory_E2E_DiscoverSuccess(t *testing.T) { + ctx := context.Background() + + server := NewServer(&Implementation{Name: "stdio-like-server", Version: "v1"}, nil) + ct, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server.Connect: %v", err) + } + defer ss.Close() + + client := NewClient(&Implementation{Name: "stdio-like-client", Version: "v1"}, nil) + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client.Connect: %v", err) + } + defer cs.Close() + + ir := cs.InitializeResult() + if ir == nil { + t.Fatal("InitializeResult is nil; discover should have populated it") + } + if ir.ProtocolVersion != protocolVersion20260728 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (negotiated via discover, no initialize)", + ir.ProtocolVersion, protocolVersion20260728) + } + if ir.ServerInfo == nil || ir.ServerInfo.Name != "stdio-like-server" { + t.Errorf("InitializeResult.ServerInfo = %+v, want name=stdio-like-server", ir.ServerInfo) + } + + // Prove the session is usable. + if _, err := cs.ListTools(ctx, nil); err != nil { + t.Errorf("ListTools after discover: %v", err) + } +} + +// TestInMemory_E2E_DiscoverFallback_NoOverlap verifies the fallback path +// over an InMemory (STDIO-equivalent) transport: the client probes with +// _meta.protocolVersion = 2026-07-28, but the server overrides discover via +// middleware to advertise only legacy versions, so the client must fall +// back to the legacy initialize handshake. +func TestInMemory_E2E_DiscoverFallback_NoOverlap(t *testing.T) { + ctx := context.Background() + server := NewServer(&Implementation{Name: "vpre-like-server", Version: "v1"}, nil) + // Intercept discover and reply as if we were a server that only + // supports legacy versions. + server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == methodDiscover { + return &DiscoverResult{ + Meta: Meta{ + MetaKeyServerInfo: &Implementation{Name: "vpre-like-server", Version: "v1"}, + }, + SupportedVersions: []string{protocolVersion20251125}, + Capabilities: &ServerCapabilities{}, + }, nil + } + return next(ctx, method, req) + } + }) + + ct, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server.Connect: %v", err) + } + defer ss.Close() + + client := NewClient(&Implementation{Name: "new-client", Version: "v1"}, nil) + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client.Connect: %v", err) + } + defer cs.Close() + + ir := cs.InitializeResult() + if ir == nil { + t.Fatal("InitializeResult is nil after fallback initialize") + } + // The fallback initialize explicitly requests protocolVersion20251125 + // (see client.go), so the server negotiates that version. + if ir.ProtocolVersion != protocolVersion20251125 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (legacy fallback after no-overlap discover)", + ir.ProtocolVersion, protocolVersion20251125) + } + + // Prove the session is usable after fallback. + if _, err := cs.ListTools(ctx, nil); err != nil { + t.Errorf("ListTools after fallback initialize: %v", err) + } +} + +// TestInMemory_E2E_DiscoverFallback_MethodNotFound verifies the fallback +// path over InMemory when the server doesn't know about server/discover at +// all (simulating a true pre-SEP-2575 server). +func TestInMemory_E2E_DiscoverFallback_MethodNotFound(t *testing.T) { + ctx := context.Background() + + server := NewServer(&Implementation{Name: "vpre-server", Version: "v1"}, nil) + server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == methodDiscover { + return nil, jsonrpc2.ErrMethodNotFound + } + return next(ctx, method, req) + } + }) + + ct, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server.Connect: %v", err) + } + defer ss.Close() + + client := NewClient(&Implementation{Name: "new-client", Version: "v1"}, nil) + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client.Connect: %v", err) + } + defer cs.Close() + + ir := cs.InitializeResult() + if ir == nil { + t.Fatal("InitializeResult is nil after fallback initialize") + } + // The fallback initialize explicitly requests protocolVersion20251125 + // (see client.go), so the server negotiates that version. + if ir.ProtocolVersion != protocolVersion20251125 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (legacy fallback after MethodNotFound)", + ir.ProtocolVersion, protocolVersion20251125) + } + + if _, err := cs.ListTools(ctx, nil); err != nil { + t.Errorf("ListTools after fallback initialize: %v", err) + } +} + +// TestInMemory_E2E_DiscoverFallback_UnsupportedProtocolVersion verifies the +// fallback path when the server explicitly rejects the discover probe with +// CodeUnsupportedProtocolVersion (the structured SEP-2575 signal). This +// exercises Path A of the fallback logic in client.go. +func TestInMemory_E2E_DiscoverFallback_UnsupportedProtocolVersion(t *testing.T) { + ctx := context.Background() + + server := NewServer(&Implementation{Name: "strict-server", Version: "v1"}, nil) + server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == methodDiscover { + return nil, &jsonrpc.Error{ + Code: CodeUnsupportedProtocolVersion, + Message: "unsupported protocol version", + } + } + return next(ctx, method, req) + } + }) + + ct, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server.Connect: %v", err) + } + defer ss.Close() + + client := NewClient(&Implementation{Name: "new-client", Version: "v1"}, nil) + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client.Connect: %v", err) + } + defer cs.Close() + + ir := cs.InitializeResult() + if ir == nil { + t.Fatal("InitializeResult is nil after fallback initialize") + } + // The fallback initialize explicitly requests protocolVersion20251125 + // (see client.go), so the server negotiates that version. + if ir.ProtocolVersion != protocolVersion20251125 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (legacy fallback after UnsupportedProtocolVersion)", + ir.ProtocolVersion, protocolVersion20251125) + } +} + +// TestClientConnectDiscover_UnsupportedVersionNegotiation verifies the +// per SEP-2575 Version Negotiation Flow: when the client probes server/discover +// with a protocol version the server doesn't implement, the server's +// UnsupportedProtocolVersionError carries Data.Supported, and the client +// selects a mutually supported version from that list and retries. +func TestClientConnectDiscover_UnsupportedVersionNegotiation(t *testing.T) { + ctx := context.Background() + + const unsupportedClientVersion = "2099-12-31" + + var ( + discoverCalls atomic.Int32 + gotInitialize atomic.Bool + observedDiscoverVersion atomic.Value // string + ) + + s := NewServer(testImpl, nil) + s.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + switch method { + case methodDiscover: + sr, ok := req.(*ServerRequest[*DiscoverParams]) + if !ok { + t.Errorf("discover req has unexpected type %T", req) + return nil, jsonrpc2.ErrMethodNotFound + } + if v, _ := sr.Params.GetMeta()[MetaKeyProtocolVersion].(string); v != "" { + observedDiscoverVersion.Store(v) + } + discoverCalls.Add(1) + case methodInitialize: + gotInitialize.Store(true) + } + return next(ctx, method, req) + } + }) + + ct, st := NewInMemoryTransports() + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server Connect: %v", err) + } + defer ss.Close() + + c := NewClient(testImpl, nil) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: unsupportedClientVersion}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer cs.Close() + + if got, want := discoverCalls.Load(), int32(1); got != want { + t.Errorf("server/discover handler call count = %d, want %d (first probe is rejected by the dispatcher; only the retry reaches the handler)", got, want) + } + if got, _ := observedDiscoverVersion.Load().(string); got != protocolVersion20260728 { + t.Errorf("retry discover requested version = %q, want %q (highest mutually supported version)", got, protocolVersion20260728) + } + if gotInitialize.Load() { + t.Error("legacy initialize handshake ran, but negotiated discover should have succeeded") + } + + ir := cs.InitializeResult() + if ir == nil { + t.Fatal("InitializeResult is nil after Connect") + } + if got, want := ir.ProtocolVersion, protocolVersion20260728; got != want { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q", got, want) + } +} diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index ebf6b59..7c98207 100644 --- a/mcp/cmd_test.go +++ b/mcp/cmd_test.go @@ -14,6 +14,7 @@ import ( "runtime" "syscall" "testing" + "testing/synctest" "time" "github.com/google/go-cmp/cmp" @@ -77,45 +78,43 @@ func runCancelContextServer() { } func TestServerRunContextCancel(t *testing.T) { - server := mcp.NewServer(&mcp.Implementation{Name: "greeter", Version: "v0.0.1"}, nil) - mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi) + synctest.Test(t, func(t *testing.T) { + server := mcp.NewServer(&mcp.Implementation{Name: "greeter", Version: "v0.0.1"}, nil) + mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() - serverTransport, clientTransport := mcp.NewInMemoryTransports() + serverTransport, clientTransport := mcp.NewInMemoryTransports() - // run the server and capture the exit error - onServerExit := make(chan error) - go func() { - onServerExit <- server.Run(ctx, serverTransport) - }() + // run the server and capture the exit error + onServerExit := make(chan error) + go func() { + onServerExit <- server.Run(ctx, serverTransport) + }() - // send a ping to the server to ensure it's running - client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) - session, err := client.Connect(ctx, clientTransport, nil) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { session.Close() }) + // send a ping to the server to ensure it's running + client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) + session, err := client.Connect(ctx, clientTransport, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { session.Close() }) - if err := session.Ping(context.Background(), nil); err != nil { - t.Fatal(err) - } + if err := session.Ping(context.Background(), nil); err != nil { + t.Fatal(err) + } - // cancel the context to stop the server - cancel() + // cancel the context to stop the server + cancel() - // wait for the server to exit - // TODO: use synctest when availble - select { - case <-time.After(5 * time.Second): - t.Fatal("server did not exit after context cancellation") - case err := <-onServerExit: + // wait for the server to exit + + err = <-onServerExit if !errors.Is(err, context.Canceled) { t.Fatalf("server did not exit after context cancellation, got error: %v", err) } - } + }) } func TestServerInterrupt(t *testing.T) { @@ -201,8 +200,7 @@ func TestStdioContextCancellation(t *testing.T) { func TestCmdTransport(t *testing.T) { requireExec(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() cmd := createServerCommand(t, "default") @@ -255,7 +253,7 @@ func TestCommandTransportTerminateDuration(t *testing.T) { requireExec(t) // Unfortunately, since it does I/O, this test needs to rely on timing (we - // can't use synctest). However, we can still decreate the default + // can't use synctest). However, we can still decrease the default // termination duration to speed up the test. const defaultDur = 50 * time.Millisecond defer mcp.SetDefaultTerminateDuration(defaultDur)() @@ -288,8 +286,7 @@ func TestCommandTransportTerminateDuration(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() // Use a command that won't exit when stdin is closed cmd := exec.Command("sleep", "20") diff --git a/mcp/conformance_go124_test.go b/mcp/conformance_go124_test.go deleted file mode 100644 index 3135970..0000000 --- a/mcp/conformance_go124_test.go +++ /dev/null @@ -1,16 +0,0 @@ -// Copyright 2025 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -//go:build go1.24 && goexperiment.synctest - -package mcp - -import ( - "testing" - "testing/synctest" -) - -func runSyncTest(t *testing.T, f func(t *testing.T)) { - synctest.Run(func() { f(t) }) -} diff --git a/mcp/conformance_go125_test.go b/mcp/conformance_go125_test.go deleted file mode 100644 index c1bb0e1..0000000 --- a/mcp/conformance_go125_test.go +++ /dev/null @@ -1,16 +0,0 @@ -// Copyright 2025 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -//go:build go1.25 - -package mcp - -import ( - "testing" - "testing/synctest" -) - -func runSyncTest(t *testing.T, f func(t *testing.T)) { - synctest.Test(t, f) -} diff --git a/mcp/conformance_test.go b/mcp/conformance_test.go index 3393efc..d7f4f74 100644 --- a/mcp/conformance_test.go +++ b/mcp/conformance_test.go @@ -2,8 +2,6 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. -//go:build (go1.24 && goexperiment.synctest) || go1.25 - package mcp import ( @@ -89,12 +87,9 @@ func TestServerConformance(t *testing.T) { // By comparison, gopls has a complicated framework based on progress // reporting and careful accounting to detect when all 'expected' work // on the server is complete. - runSyncTest(t, func(t *testing.T) { runServerTest(t, test) }) - - // TODO: in 1.25, use the following instead: - // synctest.Test(t, func(t *testing.T) { - // runServerTest(t, test) - // }) + synctest.Test(t, func(t *testing.T) { + runServerTest(t, test) + }) }) } } @@ -111,6 +106,19 @@ func structuredTool(ctx context.Context, req *CallToolRequest, args *structuredI return nil, &structuredOutput{"Ack " + args.In}, nil } +func contentTool(ctx context.Context, req *CallToolRequest, args *structuredInput) (*CallToolResult, any, error) { + return &CallToolResult{ + Content: []Content{ + &ResourceLink{ + Name: "Example Resource Link with Icons", + MIMEType: "text/plain", + URI: "https://example.com/resource/" + args.In, + Icons: []Icon{iconObj}, + }, + }, + }, nil, nil +} + type tomorrowInput struct { Now time.Time } @@ -135,12 +143,51 @@ func incTool(_ context.Context, _ *CallToolRequest, args incInput) (*CallToolRes return nil, incOutput{args.X + 1}, nil } +// confirmThenGreet exercises the multi-round-trip pattern in +// conformance fixtures. +func confirmThenGreet(_ context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &CallToolResult{ + InputRequests: InputRequestMap{ + "who": &ElicitParams{ + Message: "What is your name?", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "name": {Type: "string"}, + }, + }, + }, + }, + RequestState: "step=1", + }, nil, nil + } + name := req.Params.InputResponses["who"].(*ElicitResult).Content["name"].(string) + return &CallToolResult{ + Content: []Content{&TextContent{Text: "hello " + name}}, + }, nil, nil +} + +var iconObj = Icon{ + Source: "foobar", + MIMEType: "image/png", + Sizes: []string{"48x48", "96x96"}, + Theme: IconThemeLight, +} + // runServerTest runs the server conformance test. // It must be executed in a synctest bubble. func runServerTest(t *testing.T, test *conformanceTest) { ctx := t.Context() // Construct the server based on features listed in the test. - s := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + impl := &Implementation{Name: "testServer", Version: "v1.0.0"} + + if test.name == "spec-sep-973-additional-metadata.txtar" { + impl.Icons = []Icon{iconObj} + impl.WebsiteURL = "https://github.com/modelcontextprotocol/go-sdk" + } + + s := NewServer(impl, nil) for _, tn := range test.tools { switch tn { case "greet": @@ -148,6 +195,12 @@ func runServerTest(t *testing.T, test *conformanceTest) { Name: "greet", Description: "say hi", }, sayHi) + case "greetWithIcon": + AddTool(s, &Tool{ + Name: "greetWithIcon", + Description: "say hi", + Icons: []Icon{iconObj}, + }, sayHi) case "structured": AddTool(s, &Tool{Name: "structured"}, structuredTool) case "tomorrow": @@ -159,6 +212,17 @@ func runServerTest(t *testing.T, test *conformanceTest) { } inSchema.Properties["x"].Default = json.RawMessage(`6`) AddTool(s, &Tool{Name: "inc", InputSchema: inSchema}, incTool) + case "contentTool": + AddTool(s, &Tool{ + Name: "contentTool", + Title: "contentTool", + Description: "return resourceLink content with Icon", + }, contentTool) + case "confirmThenGreet": + AddTool(s, &Tool{ + Name: "confirmThenGreet", + Description: "ask the user for their name via elicitation, then greet them", + }, confirmThenGreet) default: t.Fatalf("unknown tool %q", tn) } @@ -167,6 +231,13 @@ func runServerTest(t *testing.T, test *conformanceTest) { switch pn { case "code_review": s.AddPrompt(codeReviewPrompt, codReviewPromptHandler) + case "code_reviewWithIcon": + s.AddPrompt(&Prompt{ + Name: "code_review", + Description: "do a code review", + Arguments: []*PromptArgument{{Name: "Code", Required: true}}, + Icons: []Icon{iconObj}, + }, codReviewPromptHandler) default: t.Fatalf("unknown prompt %q", pn) } @@ -177,6 +248,13 @@ func runServerTest(t *testing.T, test *conformanceTest) { s.AddResource(resource1, readHandler) case "info": s.AddResource(resource3, handleEmbeddedResource) + case "infoWithIcon": + s.AddResource(&Resource{ + Name: "info", + MIMEType: "text/plain", + URI: "embedded:info", + Icons: []Icon{iconObj}, + }, handleEmbeddedResource) default: t.Fatalf("unknown resource %q", rn) } @@ -275,7 +353,7 @@ func runServerTest(t *testing.T, test *conformanceTest) { // Before closing the stream, wait for all messages to be processed. synctest.Wait() if err != nil { - t.Fatalf("reading server messages failedd: %v", err) + t.Fatalf("reading server messages failed: %v", err) } if extra != nil { t.Fatalf("got extra response: %v", extra) diff --git a/mcp/content.go b/mcp/content.go index e53cad1..7af90fe 100644 --- a/mcp/content.go +++ b/mcp/content.go @@ -9,12 +9,16 @@ package mcp import ( "encoding/json" - "errors" "fmt" + + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" ) // A Content is a [TextContent], [ImageContent], [AudioContent], -// [ResourceLink], or [EmbeddedResource]. +// [ResourceLink], [EmbeddedResource], [ToolUseContent], or [ToolResultContent]. +// +// Note: [ToolUseContent] and [ToolResultContent] are only valid in sampling +// message contexts (CreateMessageParams/CreateMessageResult). type Content interface { MarshalJSON() ([]byte, error) fromWire(*wireContent) @@ -130,6 +134,8 @@ type ResourceLink struct { Size *int64 Meta Meta Annotations *Annotations + // Icons for the resource link, if any. + Icons []Icon `json:"icons,omitempty"` } func (c *ResourceLink) MarshalJSON() ([]byte, error) { @@ -143,6 +149,7 @@ func (c *ResourceLink) MarshalJSON() ([]byte, error) { Size: c.Size, Meta: c.Meta, Annotations: c.Annotations, + Icons: c.Icons, }) } @@ -155,6 +162,7 @@ func (c *ResourceLink) fromWire(wire *wireContent) { c.Size = wire.Size c.Meta = wire.Meta c.Annotations = wire.Annotations + c.Icons = wire.Icons } // EmbeddedResource contains embedded resources. @@ -179,68 +187,177 @@ func (c *EmbeddedResource) fromWire(wire *wireContent) { c.Annotations = wire.Annotations } +// ToolUseContent represents a request from the assistant to invoke a tool. +// This content type is only valid in sampling messages. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +type ToolUseContent struct { + // ID is a unique identifier for this tool use, used to match with ToolResultContent. + ID string + // Name is the name of the tool to invoke. + Name string + // Input contains the tool arguments as a JSON object. + Input map[string]any + Meta Meta +} + +func (c *ToolUseContent) MarshalJSON() ([]byte, error) { + input := c.Input + if input == nil { + input = map[string]any{} + } + wire := struct { + Type string `json:"type"` + ID string `json:"id"` + Name string `json:"name"` + Input map[string]any `json:"input"` + Meta Meta `json:"_meta,omitempty"` + }{ + Type: "tool_use", + ID: c.ID, + Name: c.Name, + Input: input, + Meta: c.Meta, + } + return json.Marshal(wire) +} + +func (c *ToolUseContent) fromWire(wire *wireContent) { + c.ID = wire.ID + c.Name = wire.Name + c.Input = wire.Input + c.Meta = wire.Meta +} + +// ToolResultContent represents the result of a tool invocation. +// This content type is only valid in sampling messages with role "user". +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +type ToolResultContent struct { + // ToolUseID references the ID from the corresponding ToolUseContent. + ToolUseID string + // Content holds the unstructured result of the tool call. + Content []Content + // StructuredContent holds an optional structured result. Per SEP-2106, it + // may be any valid JSON value (object, array, or primitive) conforming to + // the tool's output schema. + StructuredContent any + // IsError indicates whether the tool call ended in an error. + IsError bool + Meta Meta +} + +func (c *ToolResultContent) MarshalJSON() ([]byte, error) { + // Marshal nested content + var contentWire []*wireContent + for _, content := range c.Content { + data, err := content.MarshalJSON() + if err != nil { + return nil, err + } + var w wireContent + if err := internaljson.Unmarshal(data, &w); err != nil { + return nil, err + } + contentWire = append(contentWire, &w) + } + if contentWire == nil { + contentWire = []*wireContent{} // avoid JSON null + } + + wire := struct { + Type string `json:"type"` + ToolUseID string `json:"toolUseId"` + Content []*wireContent `json:"content"` + StructuredContent any `json:"structuredContent,omitempty"` + IsError bool `json:"isError,omitempty"` + Meta Meta `json:"_meta,omitempty"` + }{ + Type: "tool_result", + ToolUseID: c.ToolUseID, + Content: contentWire, + StructuredContent: c.StructuredContent, + IsError: c.IsError, + Meta: c.Meta, + } + return json.Marshal(wire) +} + +func (c *ToolResultContent) fromWire(wire *wireContent) { + c.ToolUseID = wire.ToolUseID + c.StructuredContent = wire.StructuredContent + c.IsError = wire.IsError + c.Meta = wire.Meta + // Content is handled separately in contentFromWire due to nested content +} + // ResourceContents contains the contents of a specific resource or // sub-resource. type ResourceContents struct { URI string `json:"uri"` MIMEType string `json:"mimeType,omitempty"` Text string `json:"text,omitempty"` - Blob []byte `json:"blob,omitempty"` + Blob []byte `json:"blob,omitzero"` Meta Meta `json:"_meta,omitempty"` } -func (r *ResourceContents) MarshalJSON() ([]byte, error) { - // If we could assume Go 1.24, we could use omitzero for Blob and avoid this method. - if r.URI == "" { - return nil, errors.New("ResourceContents missing URI") - } - if r.Blob == nil { - // Text. Marshal normally. - type wireResourceContents ResourceContents // (lacks MarshalJSON method) - return json.Marshal((wireResourceContents)(*r)) - } - // Blob. - if r.Text != "" { - return nil, errors.New("ResourceContents has non-zero Text and Blob fields") - } - // r.Blob may be the empty slice, so marshal with an alternative definition. - br := struct { - URI string `json:"uri,omitempty"` - MIMEType string `json:"mimeType,omitempty"` - Blob []byte `json:"blob"` - Meta Meta `json:"_meta,omitempty"` - }{ - URI: r.URI, - MIMEType: r.MIMEType, - Blob: r.Blob, - Meta: r.Meta, - } - return json.Marshal(br) -} - // wireContent is the wire format for content. // It represents the protocol types TextContent, ImageContent, AudioContent, -// ResourceLink, and EmbeddedResource. +// ResourceLink, EmbeddedResource, ToolUseContent, and ToolResultContent. // The Type field distinguishes them. In the protocol, each type has a constant // value for the field. -// At most one of Text, Data, Resource, and URI is non-zero. type wireContent struct { - Type string `json:"type"` - Text string `json:"text,omitempty"` - MIMEType string `json:"mimeType,omitempty"` - Data []byte `json:"data,omitempty"` - Resource *ResourceContents `json:"resource,omitempty"` - URI string `json:"uri,omitempty"` - Name string `json:"name,omitempty"` - Title string `json:"title,omitempty"` - Description string `json:"description,omitempty"` - Size *int64 `json:"size,omitempty"` - Meta Meta `json:"_meta,omitempty"` - Annotations *Annotations `json:"annotations,omitempty"` + Type string `json:"type"` + Text string `json:"text,omitempty"` // TextContent + MIMEType string `json:"mimeType,omitempty"` // ImageContent, AudioContent, ResourceLink + Data []byte `json:"data,omitempty"` // ImageContent, AudioContent + Resource *ResourceContents `json:"resource,omitempty"` // EmbeddedResource + URI string `json:"uri,omitempty"` // ResourceLink + Name string `json:"name,omitempty"` // ResourceLink, ToolUseContent + Title string `json:"title,omitempty"` // ResourceLink + Description string `json:"description,omitempty"` // ResourceLink + Size *int64 `json:"size,omitempty"` // ResourceLink + Meta Meta `json:"_meta,omitempty"` // all types + Annotations *Annotations `json:"annotations,omitempty"` // all types except ToolUseContent, ToolResultContent + Icons []Icon `json:"icons,omitempty"` // ResourceLink + ID string `json:"id,omitempty"` // ToolUseContent + Input map[string]any `json:"input,omitempty"` // ToolUseContent + ToolUseID string `json:"toolUseId,omitempty"` // ToolResultContent + NestedContent []*wireContent `json:"content,omitempty"` // ToolResultContent + StructuredContent any `json:"structuredContent,omitempty"` // ToolResultContent + IsError bool `json:"isError,omitempty"` // ToolResultContent +} + +// unmarshalContent unmarshals JSON that is either a single content object or +// an array of content objects. A single object is wrapped in a one-element slice. +func unmarshalContent(raw json.RawMessage, allow map[string]bool) ([]Content, error) { + if len(raw) == 0 || string(raw) == "null" { + return nil, fmt.Errorf("nil content") + } + // Try array first, then fall back to single object. + var wires []*wireContent + if err := internaljson.Unmarshal(raw, &wires); err == nil { + return contentsFromWire(wires, allow) + } + var wire wireContent + if err := internaljson.Unmarshal(raw, &wire); err != nil { + return nil, err + } + c, err := contentFromWire(&wire, allow) + if err != nil { + return nil, err + } + return []Content{c}, nil } func contentsFromWire(wires []*wireContent, allow map[string]bool) ([]Content, error) { - var blocks []Content + blocks := make([]Content, 0, len(wires)) for _, wire := range wires { block, err := contentFromWire(wire, allow) if err != nil { @@ -279,6 +396,27 @@ func contentFromWire(wire *wireContent, allow map[string]bool) (Content, error) v := new(EmbeddedResource) v.fromWire(wire) return v, nil + case "tool_use": + v := new(ToolUseContent) + v.fromWire(wire) + return v, nil + case "tool_result": + v := new(ToolResultContent) + v.fromWire(wire) + // Handle nested content - tool_result content can contain text, image, audio, + // resource_link, and resource (same as CallToolResult.content) + if wire.NestedContent != nil { + toolResultContentAllow := map[string]bool{ + "text": true, "image": true, "audio": true, + "resource_link": true, "resource": true, + } + nestedContent, err := contentsFromWire(wire.NestedContent, toolResultContentAllow) + if err != nil { + return nil, fmt.Errorf("tool_result nested content: %w", err) + } + v.Content = nestedContent + } + return v, nil } - return nil, fmt.Errorf("internal error: unrecognized content type %s", wire.Type) + return nil, fmt.Errorf("unrecognized content type %q", wire.Type) } diff --git a/mcp/content_nil_test.go b/mcp/content_nil_test.go index 8cc7bdb..f759986 100644 --- a/mcp/content_nil_test.go +++ b/mcp/content_nil_test.go @@ -15,6 +15,7 @@ import ( "testing" "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -22,8 +23,8 @@ func TestContentUnmarshalNil(t *testing.T) { tests := []struct { name string json string - content interface{} - want interface{} + content any + want any }{ { name: "CallToolResult nil Content", @@ -83,7 +84,7 @@ func TestContentUnmarshalNilWithDifferentTypes(t *testing.T) { tests := []struct { name string json string - content interface{} + content any expectError bool }{ { @@ -145,7 +146,7 @@ func TestContentUnmarshalNilWithEmptyContent(t *testing.T) { tests := []struct { name string json string - content interface{} + content any expectError bool }{ { @@ -186,7 +187,7 @@ func TestContentUnmarshalNilWithInvalidContent(t *testing.T) { tests := []struct { name string json string - content interface{} + content any expectError bool }{ { @@ -223,4 +224,12 @@ func TestContentUnmarshalNilWithInvalidContent(t *testing.T) { } } -var ctrCmpOpts = []cmp.Option{cmp.AllowUnexported(mcp.CallToolResult{})} +var ctrCmpOpts = []cmp.Option{ + cmpopts.IgnoreUnexported(mcp.CallToolResult{}, mcp.GetPromptResult{}, mcp.ReadResourceResult{}), + // Server responses under the >= 2026-07-28 protocol carry an auto-populated + // [mcp.MetaKeyServerInfo] entry; tests that compare result bodies against + // hand-crafted expected values should ignore it. + cmpopts.IgnoreFields(mcp.CallToolResult{}, "Meta"), + cmpopts.IgnoreFields(mcp.GetPromptResult{}, "Meta"), + cmpopts.IgnoreFields(mcp.ReadResourceResult{}, "Meta"), +} diff --git a/mcp/content_test.go b/mcp/content_test.go index 9366b0d..094950e 100644 --- a/mcp/content_test.go +++ b/mcp/content_test.go @@ -121,8 +121,9 @@ func TestContent(t *testing.T) { Description: "This resource demonstrates all fields", MIMEType: "text/plain", Meta: mcp.Meta{"custom": "metadata"}, + Icons: []mcp.Icon{{Source: "foobar", MIMEType: "image/png", Sizes: []string{"48x48"}, Theme: mcp.IconThemeLight}}, }, - `{"type":"resource_link","mimeType":"text/plain","uri":"https://example.com/resource","name":"Example Resource","title":"A comprehensive example resource","description":"This resource demonstrates all fields","_meta":{"custom":"metadata"}}`, + `{"type":"resource_link","mimeType":"text/plain","uri":"https://example.com/resource","name":"Example Resource","title":"A comprehensive example resource","description":"This resource demonstrates all fields","_meta":{"custom":"metadata"},"icons":[{"src":"foobar","mimeType":"image/png","sizes":["48x48"],"theme":"light"}]}`, }, } @@ -192,3 +193,53 @@ func TestEmbeddedResource(t *testing.T) { } } } + +// TestContentUnmarshal tests that unmarshaling JSON into various Content types +// works correctly, including when the Content fields are initially nil. +func TestContentUnmarshal(t *testing.T) { + valInt64 := int64(24) + tests := []struct { + name string + json string + content mcp.Content + expectContent mcp.Content + }{ + { + name: "ResourceLink", + json: `{"type":"resource_link","mimeType":"text/plain","uri":"https://example.com/resource","name":"Example Resource","title":"A comprehensive example resource","description":"This resource demonstrates all fields","_meta":{"custom":"metadata"},"icons":[{"src":"foobar","mimeType":"image/png","sizes":["48x48"],"theme":"light"}], "size":24,"annotations":{"audience":["user","assistant"],"lastModified":"2025-01-12T15:00:58Z","priority":0.5}}`, + content: &mcp.ResourceLink{}, + expectContent: &mcp.ResourceLink{ + URI: "https://example.com/resource", + Name: "Example Resource", + Title: "A comprehensive example resource", + Description: "This resource demonstrates all fields", + MIMEType: "text/plain", + // Meta: mcp.Meta{"custom": "metadata"}, + Size: &valInt64, + Annotations: &mcp.Annotations{Audience: []mcp.Role{"user", "assistant"}, LastModified: "2025-01-12T15:00:58Z", Priority: 0.5}, + Icons: []mcp.Icon{{Source: "foobar", MIMEType: "image/png", Sizes: []string{"48x48"}, Theme: mcp.IconThemeLight}}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Test that unmarshaling doesn't panic on nil Content fields + defer func() { + if r := recover(); r != nil { + t.Errorf("UnmarshalJSON panicked: %v", r) + } + }() + + err := json.Unmarshal([]byte(tt.json), tt.content) + if err != nil { + t.Errorf("UnmarshalJSON failed: %v", err) + } + + // Verify that the Content field was properly populated + if cmp.Diff(tt.expectContent, tt.content) != "" { + t.Errorf("Content is not equal: %v", cmp.Diff(tt.expectContent, tt.content)) + } + }) + } +} diff --git a/mcp/elicitation_test.go b/mcp/elicitation_test.go new file mode 100644 index 0000000..098fce3 --- /dev/null +++ b/mcp/elicitation_test.go @@ -0,0 +1,314 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "context" + "errors" + "strings" + "testing" + "testing/synctest" + + "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" +) + +// TODO: migrate other elicitation tests here. + +func TestElicitationURLMode(t *testing.T) { + ctx := context.Background() + clientErr := errors.New("client failed to elicit") + + testCases := []struct { + name string + handler func(context.Context, *ElicitRequest) (*ElicitResult, error) + params *ElicitParams + wantResultAction string + wantErrMsg string + wantErrCode int64 + }{ + { + name: "success", + handler: func(ctx context.Context, req *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + params: &ElicitParams{ + Mode: "url", + Message: "Please provide information via URL", + URL: "https://example.com/form", + }, + wantResultAction: "accept", + }, + { + name: "decline", + handler: func(ctx context.Context, req *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "decline"}, nil + }, + params: &ElicitParams{ + Mode: "url", + Message: "Please provide information via URL", + URL: "https://example.com/form", + }, + wantResultAction: "decline", + }, + { + name: "client error", + handler: func(ctx context.Context, req *ElicitRequest) (*ElicitResult, error) { + return nil, clientErr + }, + params: &ElicitParams{ + Mode: "url", + Message: "This should fail", + URL: "https://example.com/form", + }, + wantErrMsg: clientErr.Error(), + }, + { + name: "missing url", + handler: func(ctx context.Context, req *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + params: &ElicitParams{ + Mode: "url", + Message: "URL is missing", + }, + wantErrMsg: "URL must be set for URL elicitation", + wantErrCode: jsonrpc.CodeInvalidParams, + }, + { + name: "schema not allowed", + handler: func(ctx context.Context, req *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + params: &ElicitParams{ + Mode: "url", + Message: "Schema is not allowed", + URL: "https://example.com/form", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + }, + }, + wantErrMsg: "requestedSchema must not be set for URL elicitation", + wantErrCode: jsonrpc.CodeInvalidParams, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ct, st := NewInMemoryTransports() + s := NewServer(testImpl, nil) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + c := NewClient(testImpl, &ClientOptions{ + Capabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Elicitation: &ElicitationCapabilities{ + URL: &URLElicitationCapabilities{}, + }, + }, + ElicitationHandler: tc.handler, + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + result, err := ss.Elicit(ctx, tc.params) + + if tc.wantErrMsg != "" { + if err == nil || !strings.Contains(err.Error(), tc.wantErrMsg) { + t.Fatalf("Elicit(...): got error %v, want containing %q", err, tc.wantErrMsg) + } + if tc.wantErrCode != 0 { + if code := errorCode(err); code != tc.wantErrCode { + t.Errorf("Elicit(...): got error code %d, want %d", code, tc.wantErrCode) + } + } + } else { + if err != nil { + t.Fatalf("Elicit failed: %v", err) + } + if result.Action != tc.wantResultAction { + t.Errorf("Elicit(...): got action %q, want %q", result.Action, tc.wantResultAction) + } + } + }) + } +} + +func TestElicitationCompleteNotification(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + + var elicitationCompleteCh = make(chan *ElicitationCompleteParams, 1) + + c := NewClient(testImpl, &ClientOptions{ + Capabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Elicitation: &ElicitationCapabilities{ + URL: &URLElicitationCapabilities{}, + }, + }, + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + ElicitationCompleteHandler: func(_ context.Context, req *ElicitationCompleteNotificationRequest) { + elicitationCompleteCh <- req.Params + }, + }) + + ct, st := NewInMemoryTransports() + s := NewServer(testImpl, nil) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // 1. Server initiates a URL elicitation + elicitID := "testElicitationID-123" + resp, err := ss.Elicit(ctx, &ElicitParams{ + Mode: "url", + Message: "Please complete this form: ", + URL: "https://example.com/form?id=" + elicitID, + ElicitationID: elicitID, + }) + if err != nil { + t.Fatalf("Elicit failed: %v", err) + } + if resp.Action != "accept" { + t.Fatalf("Elicit action is %q, want %q", resp.Action, "accept") + } + + // 2. Server sends elicitation complete notification (simulating out-of-band completion) + err = ss.NotifyElicitationComplete(ctx, &ElicitationCompleteParams{ + ElicitationID: elicitID, + }) + if err != nil { + t.Fatalf("failed to send elicitation complete notification: %v", err) + } + + // 3. Client should receive the notification + gotParams := <-elicitationCompleteCh + if gotParams.ElicitationID != elicitID { + t.Errorf("elicitationComplete notification ID mismatch: got %q, want %q", gotParams.ElicitationID, elicitID) + } + }) +} + +func TestNotifyElicitationCompleteInvalidParams(t *testing.T) { + ctx := context.Background() + + ct, st := NewInMemoryTransports() + s := NewServer(testImpl, nil) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + c := NewClient(testImpl, nil) + cs, err := c.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // An empty elicitation ID cannot be matched to a pending elicitation by the + // client, so it is rejected rather than sent. + for _, tc := range []struct { + name string + params *ElicitationCompleteParams + }{ + {"nil params", nil}, + {"empty ID", &ElicitationCompleteParams{}}, + } { + t.Run(tc.name, func(t *testing.T) { + err := ss.NotifyElicitationComplete(ctx, tc.params) + if !errors.Is(err, jsonrpc2.ErrInvalidParams) { + t.Errorf("NotifyElicitationComplete(%v) = %v, want %v", tc.params, err, jsonrpc2.ErrInvalidParams) + } + }) + } +} + +func TestElicitationNoValidationWithoutAccept(t *testing.T) { + ctx := context.Background() + + schema := &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "test": {Type: "string"}, + }, + Required: []string{"test"}, + } + + testCases := []struct { + name string + action string + content map[string]any + wantAction string + }{ + { + name: "cancel action", + action: "cancel", + content: nil, // Empty content should be ignored + wantAction: "cancel", + }, + { + name: "decline action", + action: "decline", + content: map[string]any{}, // Empty content should be ignored + wantAction: "decline", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ct, st := NewInMemoryTransports() + s := NewServer(testImpl, nil) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + c := NewClient(testImpl, &ClientOptions{ + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: tc.action, Content: tc.content}, nil + }, + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + res, err := ss.Elicit(ctx, &ElicitParams{ + Message: "Test bug", + RequestedSchema: schema, + }) + + if err != nil { + t.Fatalf("Elicit failed: %v", err) + } + + if res.Action != tc.wantAction { + t.Errorf("Expected action %q, got %q", tc.wantAction, res.Action) + } + }) + } +} diff --git a/mcp/error_test.go b/mcp/error_test.go new file mode 100644 index 0000000..c69fc4b --- /dev/null +++ b/mcp/error_test.go @@ -0,0 +1,430 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "context" + "encoding/json" + "errors" + "strings" + "sync/atomic" + "testing" + + "github.com/modelcontextprotocol/go-sdk/jsonrpc" +) + +// TestServerErrors validates that the server returns appropriate error codes +// for various invalid requests. +func TestServerErrors(t *testing.T) { + ctx := context.Background() + + // Set up a server with tools, prompts, and resources for testing + cs, _, cleanup := basicConnection(t, func(s *Server) { + // Add a tool with required parameters + type RequiredParams struct { + Name string `json:"name" jsonschema:"the name is required"` + } + handler := func(ctx context.Context, req *CallToolRequest, args RequiredParams) (*CallToolResult, any, error) { + return &CallToolResult{ + Content: []Content{&TextContent{Text: "success"}}, + }, nil, nil + } + AddTool(s, &Tool{Name: "validate", Description: "validates params"}, handler) + + // Add a prompt + s.AddPrompt(codeReviewPrompt, codReviewPromptHandler) + + // Add a resource that returns ResourceNotFoundError + s.AddResource( + &Resource{URI: "file:///test.txt", Name: "test", MIMEType: "text/plain"}, + func(ctx context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + return nil, ResourceNotFoundError(req.Params.URI) + }, + ) + }) + defer cleanup() + + testCases := []struct { + name string + executeCall func() error + expectedCode int64 + }{ + // Note: "missing required param" is tested separately below, because + // input validation errors are returned as tool results with IsError=true + // rather than JSON-RPC errors (see #450). + { + name: "unknown tool", + executeCall: func() error { + _, err := cs.CallTool(ctx, &CallToolParams{ + Name: "nonexistent_tool", + Arguments: map[string]any{}, + }) + return err + }, + expectedCode: jsonrpc.CodeInvalidParams, + }, + { + name: "unknown prompt", + executeCall: func() error { + _, err := cs.GetPrompt(ctx, &GetPromptParams{ + Name: "nonexistent_prompt", + Arguments: map[string]string{}, + }) + return err + }, + expectedCode: jsonrpc.CodeInvalidParams, + }, + { + name: "resource not found", + executeCall: func() error { + _, err := cs.ReadResource(ctx, &ReadResourceParams{ + URI: "file:///test.txt", + }) + return err + }, + expectedCode: CodeResourceNotFound, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + err := tc.executeCall() + if err == nil { + t.Fatal("got nil error, want non-nil") + } + + var rpcErr *jsonrpc.Error + if !errors.As(err, &rpcErr) { + t.Fatalf("got error type %T, want jsonrpc.Error: %v", err, err) + } + + if rpcErr.Code != tc.expectedCode { + t.Errorf("got error code %d, want %d", rpcErr.Code, tc.expectedCode) + } + + if rpcErr.Message == "" { + t.Error("got empty error message, want non-empty") + } + }) + } +} + +// TestResourceNotFoundErrorCode verifies ResourceNotFoundError code and +// CodeResourceNotFound are -32602 (InvalidParams) per SEP-2164. +func TestResourceNotFoundErrorCode(t *testing.T) { + err := ResourceNotFoundError("file:///test.txt") + var rpcErr *jsonrpc.Error + if !errors.As(err, &rpcErr) { + t.Fatalf("got error type %T, want *jsonrpc.Error", err) + } + if rpcErr.Code != jsonrpc.CodeInvalidParams { + t.Errorf("got error code %d, want %d (InvalidParams)", rpcErr.Code, jsonrpc.CodeInvalidParams) + } + if CodeResourceNotFound != jsonrpc.CodeInvalidParams { + t.Errorf("CodeResourceNotFound = %d, want %d", CodeResourceNotFound, jsonrpc.CodeInvalidParams) + } +} + +// TestReadResourceNilResult verifies that a resource handler returning +// (nil, nil) yields a clean error to the client instead of a nil pointer +// panic that crashes the server. The sibling handlers callTool and getPrompt +// already tolerate a (nil, nil) return; readResource must do the same. +func TestReadResourceNilResult(t *testing.T) { + ctx := context.Background() + + cs, _, cleanup := basicConnection(t, func(s *Server) { + s.AddResource( + &Resource{URI: "file:///nil.txt", Name: "nil", MIMEType: "text/plain"}, + func(ctx context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + return nil, nil + }, + ) + }) + defer cleanup() + + _, err := cs.ReadResource(ctx, &ReadResourceParams{URI: "file:///nil.txt"}) + if err == nil { + t.Fatal("got nil error, want non-nil error for nil resource result") + } + if !strings.Contains(err.Error(), "nil information") { + t.Errorf("got error %q, want it to mention 'nil information'", err.Error()) + } +} + +// TestInputValidationToolError validates that input validation errors (missing +// required params, wrong types) are returned as tool results with IsError=true, +// not as JSON-RPC errors. This allows LLMs to see the error and self-correct. +// See #450. +func TestInputValidationToolError(t *testing.T) { + ctx := context.Background() + + type RequiredParams struct { + Name string `json:"name" jsonschema:"the name is required"` + } + handler := func(ctx context.Context, req *CallToolRequest, args RequiredParams) (*CallToolResult, any, error) { + return &CallToolResult{ + Content: []Content{&TextContent{Text: "success"}}, + }, nil, nil + } + + cs, _, cleanup := basicConnection(t, func(s *Server) { + AddTool(s, &Tool{Name: "validate", Description: "validates params"}, handler) + }) + defer cleanup() + + // Call the tool with missing required "name" field. + result, err := cs.CallTool(ctx, &CallToolParams{ + Name: "validate", + Arguments: map[string]any{}, + }) + if err != nil { + t.Fatalf("CallTool returned error: %v; want tool result with IsError", err) + } + if !result.IsError { + t.Fatal("got IsError=false, want IsError=true for missing required param") + } + text := result.Content[0].(*TextContent).Text + if !strings.Contains(text, "name") { + t.Errorf("error text %q does not mention missing field \"name\"", text) + } +} + +// TestURLElicitationRequired validates that URL elicitation required errors +// are properly created and handled by the client. +func TestURLElicitationRequired(t *testing.T) { + ctx := context.Background() + + t.Run("error creation", func(t *testing.T) { + elicitations := []*ElicitParams{ + { + Mode: "url", + Message: "Please authorize", + URL: "https://example.com/auth", + ElicitationID: "auth-123", + }, + } + + err := URLElicitationRequiredError(elicitations) + + var rpcErr *jsonrpc.Error + if !errors.As(err, &rpcErr) { + t.Fatalf("got error type %T, want jsonrpc.Error", err) + } + + if rpcErr.Code != CodeURLElicitationRequired { + t.Errorf("got error code %d, want %d", rpcErr.Code, CodeURLElicitationRequired) + } + + if rpcErr.Message != "URL elicitation required" { + t.Errorf("got message %q, want 'URL elicitation required'", rpcErr.Message) + } + + if rpcErr.Data == nil { + t.Fatal("got nil error data, want non-nil") + } + + // Verify the elicitations can be unmarshaled from the error data + var errorData struct { + Elicitations []*ElicitParams `json:"elicitations"` + } + if err := json.Unmarshal(rpcErr.Data, &errorData); err != nil { + t.Fatalf("failed to unmarshal error data: %v", err) + } + + if len(errorData.Elicitations) != 1 { + t.Fatalf("got %d elicitations, want 1", len(errorData.Elicitations)) + } + + if errorData.Elicitations[0].URL != "https://example.com/auth" { + t.Errorf("got URL %q, want 'https://example.com/auth'", errorData.Elicitations[0].URL) + } + }) + + t.Run("error creation with non-URL mode panics", func(t *testing.T) { + defer func() { + if r := recover(); r == nil { + t.Error("got no panic when creating URLElicitationRequiredError with non-URL mode, want panic") + } + }() + + // This should panic because mode is "form" + URLElicitationRequiredError([]*ElicitParams{ + { + Mode: "form", + Message: "This should panic", + ElicitationID: "bad-123", + }, + }) + }) + + t.Run("error creation with empty mode panics", func(t *testing.T) { + defer func() { + if r := recover(); r == nil { + t.Error("got no panic when creating URLElicitationRequiredError with empty mode (defaults to form), want panic") + } + }() + + // This should panic because empty mode defaults to "form" + URLElicitationRequiredError([]*ElicitParams{ + { + Message: "This should panic", + ElicitationID: "bad-123", + }, + }) + }) + + t.Run("client middleware", func(t *testing.T) { + // Declare ss outside so it can be captured in handlers. + var ss *ServerSession + + elicitCalled := false + elicitURL := "" + elicitID := "form-123" + + // Create client with elicitation handler and middleware. + client := NewClient(testImpl, &ClientOptions{ + Capabilities: &ClientCapabilities{ + Roots: RootCapabilities{ListChanged: true}, + RootsV2: &RootCapabilities{ListChanged: true}, + Elicitation: &ElicitationCapabilities{ + URL: &URLElicitationCapabilities{}, + }, + }, + ElicitationHandler: func(ctx context.Context, req *ElicitRequest) (*ElicitResult, error) { + elicitCalled = true + elicitURL = req.Params.URL + + // Simulate the server sending elicitation complete notification. + // In a real scenario, this would happen out-of-band after the user + // completes the form submission. + go func() { + err := ss.NotifyElicitationComplete(ctx, &ElicitationCompleteParams{ + ElicitationID: elicitID, + }) + if err != nil { + t.Errorf("failed to send elicitation complete notification: %v", err) + } + }() + + return &ElicitResult{Action: "accept"}, nil + }, + }) + // Add URL elicitation middleware for automatic retry. + client.AddSendingMiddleware(urlElicitationMiddleware()) + + var callCount atomic.Int32 + + cs, serverSession, cleanup := basicClientServerConnection(t, + client, + nil, + func(s *Server) { + // Tool that requires form submission on first call, succeeds on second. + handler := func(ctx context.Context, req *CallToolRequest, args map[string]any) (*CallToolResult, any, error) { + if callCount.Add(1) == 1 { + // First call: require elicitation. + return nil, nil, URLElicitationRequiredError([]*ElicitParams{ + { + Mode: "url", + Message: "Please complete the form", + URL: "https://example.com/form", + ElicitationID: elicitID, + }, + }) + } + // Second call (after retry): return success. + return &CallToolResult{ + Content: []Content{&TextContent{Text: "form submitted"}}, + }, nil, nil + } + AddTool(s, &Tool{Name: "submit_form", Description: "requires form submission"}, handler) + + // Tool that returns invalid elicitation mode (form instead of URL). + badHandler := func(ctx context.Context, req *CallToolRequest, args map[string]any) (*CallToolResult, any, error) { + // Manually construct an error with form mode (bypassing validation). + data, _ := json.Marshal(map[string]any{ + "elicitations": []*ElicitParams{ + { + Mode: "form", + Message: "Invalid mode", + ElicitationID: "bad-form", + }, + }, + }) + return nil, nil, &jsonrpc.Error{ + Code: CodeURLElicitationRequired, + Message: "URL elicitation required", + Data: json.RawMessage(data), + } + } + AddTool(s, &Tool{Name: "bad_tool", Description: "returns invalid elicitation"}, badHandler) + }, + ) + ss = serverSession + defer cleanup() + + t.Run("auto-retry after elicitation", func(t *testing.T) { + // Reset state for this subtest. + elicitCalled = false + elicitURL = "" + callCount.Store(0) + + // Call the tool that requires URL elicitation. + result, err := cs.CallTool(ctx, &CallToolParams{ + Name: "submit_form", + Arguments: map[string]any{}, + }) + + // After automatic retry, the operation should succeed. + if err != nil { + t.Fatalf("CallTool failed: %v", err) + } + + // Verify the elicitation handler was called. + if !elicitCalled { + t.Error("elicitation handler not called") + } + + if elicitURL != "https://example.com/form" { + t.Errorf("got elicit URL %q, want 'https://example.com/form'", elicitURL) + } + + // Verify the tool was called twice (first attempt + retry). + if got, want := callCount.Load(), int32(2); got != want { + t.Errorf("CallTool(): with retry, got %d tool calls, want %d", got, want) + } + + // Verify we got the successful result. + if len(result.Content) != 1 { + t.Fatalf("CallTool(): got %d content items, want 1", len(result.Content)) + } + + textContent, ok := result.Content[0].(*TextContent) + if !ok { + t.Fatalf("CallTool(): got content type %T, want TextContent", result.Content[0]) + } + + if textContent.Text != "form submitted" { + t.Errorf("CallTool(): got text %q, want 'form submitted'", textContent.Text) + } + }) + + t.Run("reject non-URL mode", func(t *testing.T) { + // Call the tool that returns invalid elicitation mode. + _, err := cs.CallTool(ctx, &CallToolParams{ + Name: "bad_tool", + Arguments: map[string]any{}, + }) + + // Should get an error about invalid mode. + if err == nil { + t.Fatal("got nil error for non-URL mode elicitation, want error") + } + + if !strings.Contains(err.Error(), "URL mode") { + t.Errorf("got error %v, want mention of URL mode", err) + } + }) + }) +} diff --git a/mcp/event.go b/mcp/event.go index 281f592..f0cbb2f 100644 --- a/mcp/event.go +++ b/mcp/event.go @@ -29,18 +29,19 @@ const validateMemoryEventStore = false // An Event is a server-sent event. // See https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#fields. type Event struct { - Name string // the "event" field - ID string // the "id" field - Data []byte // the "data" field + Name string // the "event" field + ID string // the "id" field + Data []byte // the "data" field + Retry string // the "retry" field } // Empty reports whether the Event is empty. func (e Event) Empty() bool { - return e.Name == "" && e.ID == "" && len(e.Data) == 0 + return e.Name == "" && e.ID == "" && len(e.Data) == 0 && e.Retry == "" } // writeEvent writes the event to w, and flushes. -func writeEvent(w io.Writer, evt Event) (int, error) { +func writeEvent(w http.ResponseWriter, evt Event) (int, error) { var b bytes.Buffer if evt.Name != "" { fmt.Fprintf(&b, "event: %s\n", evt.Name) @@ -48,24 +49,38 @@ func writeEvent(w io.Writer, evt Event) (int, error) { if evt.ID != "" { fmt.Fprintf(&b, "id: %s\n", evt.ID) } - fmt.Fprintf(&b, "data: %s\n\n", string(evt.Data)) - n, err := w.Write(b.Bytes()) - if f, ok := w.(http.Flusher); ok { - f.Flush() + if evt.Retry != "" { + fmt.Fprintf(&b, "retry: %s\n", evt.Retry) } + // Write the payload directly into a pre-grown buffer to avoid the extra + // copy from string(evt.Data) and repeated buffer regrowths. + b.Grow(len("data: \n\n") + len(evt.Data)) + b.WriteString("data: ") + b.Write(evt.Data) + b.WriteString("\n\n") + n, err := w.Write(b.Bytes()) + rc := http.NewResponseController(w) + // Ignore returned error as flushing is best-effort. + _ = rc.Flush() return n, err } -// scanEvents iterates SSE events in the given scanner. The iterated error is -// terminal: if encountered, the stream is corrupt or broken and should no -// longer be used. +// DefaultMaxEventSize is the default maximum number of bytes buffered while +// reading a single server-sent event before the input is rejected. +const DefaultMaxEventSize = 16 << 20 // 16 MiB + +// scanEventsLimited iterates SSE events in the given reader, buffering at most +// maxEventSize bytes per event. When maxEventSize > 0, an event is reported as +// [errMalformedEvent] once its size exceeds it; a non-positive maxEventSize +// disables the cap. +// +// The iterated error is terminal: if encountered, the stream is corrupt or +// broken and should no longer be used. // // TODO(rfindley): consider a different API here that makes failure modes more // apparent. -func scanEvents(r io.Reader) iter.Seq2[Event, error] { - scanner := bufio.NewScanner(r) - const maxTokenSize = 1 * 1024 * 1024 // 1 MiB max line size - scanner.Buffer(nil, maxTokenSize) +func scanEventsLimited(r io.Reader, maxEventSize int) iter.Seq2[Event, error] { + reader := bufio.NewReader(r) // TODO: investigate proper behavior when events are out of order, or have // non-standard names. @@ -73,6 +88,7 @@ func scanEvents(r io.Reader) iter.Seq2[Event, error] { eventKey = []byte("event") idKey = []byte("id") dataKey = []byte("data") + retryKey = []byte("retry") ) return func(yield func(Event, error) bool) { @@ -86,62 +102,79 @@ func scanEvents(r io.Reader) iter.Seq2[Event, error] { // - Lines starting with ":" are ignored. // - Records are terminated with two consecutive newlines. var ( - evt Event - dataBuf *bytes.Buffer // if non-nil, preceding field was also data + evt Event + dataBuf *bytes.Buffer // if non-nil, preceding field was also data + eventBytes int // bytes read for the current, in-progress event ) - flushData := func() { + yieldEvent := func() bool { if dataBuf != nil { evt.Data = dataBuf.Bytes() dataBuf = nil } + if evt.Empty() { + return true + } + if !yield(evt, nil) { + return false + } + evt = Event{} + return true } - for scanner.Scan() { - line := scanner.Bytes() + for { + budget := -1 + if maxEventSize > 0 { + budget = maxEventSize - eventBytes + } + line, err := readEventLine(reader, budget) + if errors.Is(err, errEventTooLarge) { + yield(Event{}, fmt.Errorf("%w: SSE event exceeded %d bytes without terminating", errMalformedEvent, maxEventSize)) + return + } + if err != nil && !errors.Is(err, io.EOF) { + yield(Event{}, fmt.Errorf("error reading event: %v", err)) + return + } + eventBytes += len(line) + line = bytes.TrimRight(line, "\r\n") + isEOF := errors.Is(err, io.EOF) + if len(line) == 0 { - flushData() - // \n\n is the record delimiter - if !evt.Empty() && !yield(evt, nil) { + if !yieldEvent() { + return + } + eventBytes = 0 // reset the budget between events + if isEOF { return } - evt = Event{} continue } before, after, found := bytes.Cut(line, []byte{':'}) if !found { - yield(Event{}, fmt.Errorf("malformed line in SSE stream: %q", string(line))) + yield(Event{}, fmt.Errorf("%w: malformed line in SSE stream: %q", errMalformedEvent, string(line))) return } - if !bytes.Equal(before, dataKey) { - flushData() - } switch { case bytes.Equal(before, eventKey): evt.Name = strings.TrimSpace(string(after)) case bytes.Equal(before, idKey): evt.ID = strings.TrimSpace(string(after)) + case bytes.Equal(before, retryKey): + evt.Retry = strings.TrimSpace(string(after)) case bytes.Equal(before, dataKey): data := bytes.TrimSpace(after) - if dataBuf != nil { - dataBuf.WriteByte('\n') - dataBuf.Write(data) - } else { + if dataBuf == nil { dataBuf = new(bytes.Buffer) - dataBuf.Write(data) + } else { + dataBuf.WriteByte('\n') } + dataBuf.Write(data) } - } - if err := scanner.Err(); err != nil { - if errors.Is(err, bufio.ErrTooLong) { - err = fmt.Errorf("event exceeded max line length of %d", maxTokenSize) - } - if !yield(Event{}, err) { + + if isEOF { + yieldEvent() return } } - flushData() - if !evt.Empty() { - yield(evt, nil) - } } } @@ -191,12 +224,8 @@ type dataList struct { } func (dl *dataList) appendData(d []byte) { - // If we allowed empty data, we would consume memory without incrementing the size. - // We could of course account for that, but we keep it simple and assume there is no - // empty data. - if len(d) == 0 { - panic("empty data item") - } + // Empty data consumes memory but doesn't increment size. However, it should + // be rare. dl.data = append(dl.data, d) dl.size += len(d) } @@ -307,6 +336,36 @@ func (s *MemoryEventStore) Append(_ context.Context, sessionID, streamID string, // index is no longer available. var ErrEventsPurged = errors.New("data purged") +// errMalformedEvent is returned when an SSE event cannot be parsed due to format violations. +// This is a hard error indicating corrupted data or protocol violations, as opposed to +// transient I/O errors which may be retryable. +var errMalformedEvent = errors.New("malformed event") + +// errEventTooLarge is returned by [readEventLine] when a single line would +// exceed the remaining per-event byte budget. +var errEventTooLarge = errors.New("SSE event exceeded maximum size") + +// readEventLine reads a single '\n'-terminated line from r. When budget >= 0 it +// reads at most budget bytes, returning [errEventTooLarge] once that budget is +// exceeded before a newline arrives. +func readEventLine(r *bufio.Reader, budget int) ([]byte, error) { + if budget < 0 { + return r.ReadBytes('\n') + } + var line []byte + for { + frag, err := r.ReadSlice('\n') + if len(line)+len(frag) > budget { + return nil, errEventTooLarge + } + line = append(line, frag...) + if errors.Is(err, bufio.ErrBufferFull) { + continue // looking for delim + } + return line, err + } +} + // After implements [EventStore.After]. func (s *MemoryEventStore) After(_ context.Context, sessionID, streamID string, index int) iter.Seq2[[]byte, error] { // Return the data items to yield. @@ -322,12 +381,15 @@ func (s *MemoryEventStore) After(_ context.Context, sessionID, streamID string, if !ok { return nil, fmt.Errorf("MemoryEventStore.After: unknown stream ID %v in session %q", streamID, sessionID) } - start := index + 1 - if dl.first > start { + start := (index + 1) - dl.first + if start < 0 { return nil, fmt.Errorf("MemoryEventStore.After: index %d, stream ID %v, session %q: %w", index, streamID, sessionID, ErrEventsPurged) } - return slices.Clone(dl.data[start-dl.first:]), nil + if start >= len(dl.data) { + return nil, nil + } + return slices.Clone(dl.data[start:]), nil } return func(yield func([]byte, error) bool) { @@ -366,10 +428,11 @@ func (s *MemoryEventStore) purge() { for _, dl := range sm { if dl.size > 0 { r := dl.removeFirst() - if r > 0 { - changed = true - s.nBytes -= r - } + // Even if we remove an empty chunk, that's + // still progress. There may be non-empty + // chunks after it. + changed = true + s.nBytes -= r } } } diff --git a/mcp/event_test.go b/mcp/event_test.go index dacf30e..d326415 100644 --- a/mcp/event_test.go +++ b/mcp/event_test.go @@ -5,8 +5,12 @@ package mcp import ( + "bytes" "context" + "crypto/rand" + "errors" "fmt" + "io" "slices" "strings" "testing" @@ -54,6 +58,78 @@ func TestScanEvents(t *testing.T) { input: "invalid line\n\n", wantErr: "malformed line", }, + { + name: "message with 2 data lines and another event", + input: "event: message\ndata: hello\ndata: hello\ndata: hello\n\nevent:keepalive", + want: []Event{ + {Name: "message", Data: []byte("hello\nhello\nhello")}, + {Name: "keepalive"}, + }, + }, + { + name: "event with multiple lines", + input: "event: message\ndata: hello\ndata: hello\ndata: hello\nid:1", + want: []Event{ + {Name: "message", ID: "1", Data: []byte("hello\nhello\nhello")}, + }, + }, + { + name: "multiple events, out of order keys", + input: strings.Join([]string{ + "event:message", + "data: hello0", + "\n", + "data: hello1", + "data: hello1", + "id:1", + "event:message", + "\n", + "event:message", + "data: hello3", + "data: hello3", + "id:3", + "\n", + "data: hello4", + "data: hello4", + "id:4", + "event:message", + }, "\n"), + want: []Event{ + {Name: "message", Data: []byte("hello0")}, + {Name: "message", ID: "1", Data: []byte("hello1\nhello1")}, + {Name: "message", ID: "3", Data: []byte("hello3\nhello3")}, + {Name: "message", ID: "4", Data: []byte("hello4\nhello4")}, + }, + }, + { + name: "non-continuous data items in the event", + input: "event: foo\ndata: 123\nretry: 5\ndata: 456", + want: []Event{ + {Name: "foo", Data: []byte("123\n456"), Retry: "5"}, + }, + }, + { + name: "no-data events", + input: "event: foo\n\nevent: bar", + want: []Event{ + {Name: "foo"}, + {Name: "bar"}, + }, + }, + { + name: "empty data event", + input: "event: foo\ndata:\n\nevent: bar", + want: []Event{ + {Name: "foo"}, + {Name: "bar"}, + }, + }, + { + + name: "malformed data event", + input: "someline", + wantErr: "malformed event", + }, } for _, tt := range tests { @@ -61,7 +137,7 @@ func TestScanEvents(t *testing.T) { r := strings.NewReader(tt.input) var got []Event var err error - for e, err2 := range scanEvents(r) { + for e, err2 := range scanEventsLimited(r, DefaultMaxEventSize) { if err2 != nil { err = err2 break @@ -102,6 +178,136 @@ func TestScanEvents(t *testing.T) { } } +// endlessReader streams a fixed prefix once, then repeats fill forever. +type endlessReader struct { + prefix string + sent bool + repeat string + repIndex int +} + +func (r *endlessReader) Read(p []byte) (int, error) { + if !r.sent { + n := copy(p, r.prefix) + r.sent = true + return n, nil + } + for i := range p { + p[i] = r.repeat[r.repIndex%len(r.repeat)] + r.repIndex++ + } + return len(p), nil +} + +func (r *endlessReader) Close() error { return nil } + +// TestScanEventsMaxEventSize verifies that scanEvents bounds the bytes +// buffered for a single event. +func TestScanEventsMaxEventSize(t *testing.T) { + wantByte := byte('A') + eventOfSize := func(size int) []byte { + var buf bytes.Buffer + buf.WriteString("data: ") + buf.Write(bytes.Repeat([]byte{wantByte}, size)) + buf.WriteString("\n\n") + return buf.Bytes() + } + + tests := []struct { + name string + reader io.Reader + maxEventSize int + wantErr bool + wantDataLengths []int + }{ + { + name: "negative budget disables the cap", + reader: bytes.NewReader(eventOfSize(512)), + maxEventSize: -1, + wantDataLengths: []int{512}, + }, + { + name: "long line under cap", // bufio.Reader default buf size is 1 << 12 + reader: bytes.NewReader(eventOfSize(1 << 14)), + maxEventSize: 1 << 15, + wantDataLengths: []int{1 << 14}, + }, + { + name: "unbounded rejected", + reader: &endlessReader{prefix: "data: ", repeat: string(wantByte)}, + maxEventSize: 1024, + wantErr: true, + }, + + { + name: "consecutive data fields longer than limit", + reader: strings.NewReader(func() string { + var b strings.Builder + for range 2048 { + b.WriteString("data: \n") + } + return b.String() + }()), + maxEventSize: 1024, + wantErr: true, + }, + { + name: "budget resets between events", + reader: bytes.NewReader(append(eventOfSize(3*1024), eventOfSize(3*1024)...)), + maxEventSize: 4 * 1024, + wantDataLengths: []int{3 * 1024, 3 * 1024}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + type result struct { + events []Event + err error + } + + done := make(chan result, 1) + go func() { + var events []Event + for e, err := range scanEventsLimited(tt.reader, tt.maxEventSize) { + if err != nil { + done <- result{events, err} + return + } + events = append(events, e) + } + done <- result{events, nil} + }() + + var res result + select { + case res = <-done: + case <-time.After(5 * time.Second): + t.Fatal("scanEvents did not return") + } + + if tt.wantErr { + if !errors.Is(res.err, errMalformedEvent) { + t.Fatalf("got error %v, want errMalformedEvent", res.err) + } + return + } + if res.err != nil { + t.Fatalf("unexpected error: %v", res.err) + } + if len(res.events) != len(tt.wantDataLengths) { + t.Fatalf("got %d events, want %d", len(res.events), len(tt.wantDataLengths)) + } + for i, wantLen := range tt.wantDataLengths { + wantRepeated := string(wantByte) + if got := res.events[i].Data; string(got) != strings.Repeat(wantRepeated, wantLen) { + t.Errorf("event %d: got %d data bytes, want %d bytes of %s", i, len(got), wantLen, wantRepeated) + } + } + }) + } +} + func TestMemoryEventStoreState(t *testing.T) { ctx := context.Background() @@ -226,7 +432,11 @@ func TestMemoryEventStoreAfter(t *testing.T) { {"S1", "1", 0, []string{"d2", "d3"}, ""}, {"S1", "1", 1, []string{"d3"}, ""}, {"S1", "1", 2, nil, ""}, + {"S1", "1", 3, nil, ""}, + {"S1", "1", 100, nil, ""}, // past the end of stream "1" (highest is 2) {"S1", "2", 0, nil, ""}, + {"S1", "2", 1, nil, ""}, + {"S1", "2", 100, nil, ""}, {"S1", "3", 0, nil, "unknown stream ID"}, {"S2", "0", 0, nil, "unknown session ID"}, } { @@ -254,6 +464,29 @@ func TestMemoryEventStoreAfter(t *testing.T) { } } +func TestMemoryEventStorePurgeEmpty(t *testing.T) { + ctx := context.Background() + s := NewMemoryEventStore(nil) + s.SetMaxBytes(100) + + // Append an empty chunk first. + if err := s.Append(ctx, "S1", "1", []byte("")); err != nil { + t.Fatal(err) + } + // Append a non-empty chunk. + if err := s.Append(ctx, "S1", "1", []byte("1234567890")); err != nil { + t.Fatal(err) + } + // Now nBytes = 10, but the first chunk is empty. + + // This should not panic. + s.SetMaxBytes(5) + + if s.nBytes > 5 { + t.Errorf("got nBytes %d, want <= 5", s.nBytes) + } +} + func BenchmarkMemoryEventStore(b *testing.B) { // Benchmark with various settings for event store size, number of session, // and payload size. @@ -281,7 +514,7 @@ func BenchmarkMemoryEventStore(b *testing.B) { for i := range sessionIDs { sessionIDs[i] = fmt.Sprint(i) for j := range 3 { - streamIDs[i][j] = randText() + streamIDs[i][j] = rand.Text() } } payload := make([]byte, test.datasize) diff --git a/mcp/logging.go b/mcp/logging.go index 208427e..e8d0cd7 100644 --- a/mcp/logging.go +++ b/mcp/logging.go @@ -10,8 +10,11 @@ import ( "context" "encoding/json" "log/slog" + "slices" "sync" "time" + + "golang.org/x/time/rate" ) // Logging levels. @@ -65,7 +68,19 @@ func compareLevels(l1, l2 LoggingLevel) int { return cmp.Compare(mcpLevelToSlog(l1), mcpLevelToSlog(l2)) } +type logLevelContextKey struct{} + +func logLevelFromContext(ctx context.Context) (LoggingLevel, bool) { + v, ok := ctx.Value(logLevelContextKey{}).(LoggingLevel) + return v, ok +} + // LoggingHandlerOptions are options for a LoggingHandler. +// +// Deprecated: the logging feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type LoggingHandlerOptions struct { // The value for the "logger" field of logging notifications. LoggerName string @@ -76,37 +91,38 @@ type LoggingHandlerOptions struct { } // A LoggingHandler is a [slog.Handler] for MCP. +// +// Deprecated: the logging feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type LoggingHandler struct { opts LoggingHandlerOptions ss *ServerSession // Ensures that the buffer reset is atomic with the write (see Handle). // A pointer so that clones share the mutex. See // https://github.com/golang/example/blob/master/slog-handler-guide/README.md#getting-the-mutex-right. - mu *sync.Mutex - lastMessageSent time.Time // for rate-limiting - buf *bytes.Buffer - handler slog.Handler + mu *sync.Mutex + limiter *rate.Limiter // for rate-limiting + buf *bytes.Buffer + handler slog.Handler } -// discardHandler is a slog.Handler that drops all logs. -// TODO: use slog.DiscardHandler when we require Go 1.24+. -type discardHandler struct{} - -func (discardHandler) Enabled(context.Context, slog.Level) bool { return false } -func (discardHandler) Handle(context.Context, slog.Record) error { return nil } -func (discardHandler) WithAttrs([]slog.Attr) slog.Handler { return discardHandler{} } -func (discardHandler) WithGroup(string) slog.Handler { return discardHandler{} } - // ensureLogger returns l if non-nil, otherwise a discard logger. func ensureLogger(l *slog.Logger) *slog.Logger { if l != nil { return l } - return slog.New(discardHandler{}) + return slog.New(slog.DiscardHandler) } // NewLoggingHandler creates a [LoggingHandler] that logs to the given [ServerSession] using a // [slog.JSONHandler]. +// +// Deprecated: the logging feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. func NewLoggingHandler(ss *ServerSession, opts *LoggingHandlerOptions) *LoggingHandler { var buf bytes.Buffer jsonHandler := slog.NewJSONHandler(&buf, &slog.HandlerOptions{ @@ -126,6 +142,9 @@ func NewLoggingHandler(ss *ServerSession, opts *LoggingHandlerOptions) *LoggingH } if opts != nil { lh.opts = *opts + if opts.MinInterval > 0 { + lh.limiter = rate.NewLimiter(rate.Every(opts.MinInterval), 1) + } } return lh } @@ -134,6 +153,9 @@ func NewLoggingHandler(ss *ServerSession, opts *LoggingHandlerOptions) *LoggingH func (h *LoggingHandler) Enabled(ctx context.Context, level slog.Level) bool { // This is also checked in ServerSession.LoggingMessage, so checking it here // is just an optimization that skips building the JSON. + if mcpLevel, ok := logLevelFromContext(ctx); ok { + return mcpLevel != "" && level >= mcpLevelToSlog(mcpLevel) + } h.ss.mu.Lock() mcpLevel := h.ss.state.LogLevel h.ss.mu.Unlock() @@ -165,16 +187,12 @@ func (h *LoggingHandler) Handle(ctx context.Context, r slog.Record) error { func (h *LoggingHandler) handle(ctx context.Context, r slog.Record) error { // Observe the rate limit. - // TODO(jba): use golang.org/x/time/rate. (We can't here because it would require adding - // golang.org/x/time to the go.mod file.) - h.mu.Lock() - skip := time.Since(h.lastMessageSent) < h.opts.MinInterval - h.mu.Unlock() - if skip { + if h.limiter != nil && !h.limiter.Allow() { return nil } var err error + var data json.RawMessage // Make the buffer reset atomic with the record write. // We are careful here in the unlikely event that the handler panics. // We don't want to hold the lock for the entire function, because Notify is @@ -185,19 +203,17 @@ func (h *LoggingHandler) handle(ctx context.Context, r slog.Record) error { defer h.mu.Unlock() h.buf.Reset() err = h.handler.Handle(ctx, r) + // Clone the buffer as Bytes() references the internal buffer. + data = json.RawMessage(slices.Clone(h.buf.Bytes())) }() if err != nil { return err } - h.mu.Lock() - h.lastMessageSent = time.Now() - h.mu.Unlock() - params := &LoggingMessageParams{ Logger: h.opts.LoggerName, Level: slogLevelToMCP(r.Level), - Data: json.RawMessage(h.buf.Bytes()), + Data: data, } // We pass the argument context to Notify, even though slog.Handler.Handle's // documentation says not to. diff --git a/mcp/mcp_example_test.go b/mcp/mcp_example_test.go index 25f39fb..120789c 100644 --- a/mcp/mcp_example_test.go +++ b/mcp/mcp_example_test.go @@ -19,19 +19,13 @@ import ( func Example_lifecycle() { ctx := context.Background() - // Create a client and server. - // Wait for the client to initialize the session. + // Create a client and server. Under protocol version 2026-07-28, there + // is no dedicated initialize/initialized handshake: the session is + // implicitly live the moment the client issues its first request. client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) - server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, &mcp.ServerOptions{ - InitializedHandler: func(context.Context, *mcp.InitializedRequest) { - fmt.Println("initialized!") - }, - }) + server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) // Connect the server and client using in-memory transports. - // - // Connect the server first so that it's ready to receive initialization - // messages from the client. t1, t2 := mcp.NewInMemoryTransports() serverSession, err := server.Connect(ctx, t1, nil) if err != nil { @@ -50,7 +44,7 @@ func Example_lifecycle() { if err := serverSession.Wait(); err != nil { log.Fatal(err) } - // Output: initialized! + // Output: } // !-lifecycle diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index a78c152..7b10b30 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -2,6 +2,8 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. +//lint:file-ignore SA1019 tests exercise deprecated SEP-2577 APIs (roots, sampling, logging). + package mcp import ( @@ -12,6 +14,8 @@ import ( "fmt" "io" "log/slog" + "net/http" + "net/http/httptest" "net/url" "path/filepath" "runtime" @@ -20,11 +24,14 @@ import ( "sync" "sync/atomic" "testing" + "testing/synctest" "time" "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" "github.com/google/jsonschema-go/jsonschema" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) type hiParams struct { @@ -56,467 +63,502 @@ func codReviewPromptHandler(_ context.Context, req *GetPromptRequest) (*GetPromp }, nil } +// TODO(maciej-kisiel): split this test into multiple test cases that target specific functionality. func TestEndToEnd(t *testing.T) { - ctx := context.Background() - var ct, st Transport = NewInMemoryTransports() + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + var ct, st Transport = NewInMemoryTransports() - // Channels to check if notification callbacks happened. - notificationChans := map[string]chan int{} - for _, name := range []string{"initialized", "roots", "tools", "prompts", "resources", "progress_server", "progress_client", "resource_updated", "subscribe", "unsubscribe"} { - notificationChans[name] = make(chan int, 1) - } - waitForNotification := func(t *testing.T, name string) { - t.Helper() - select { - case <-notificationChans[name]: - case <-time.After(time.Second): - t.Fatalf("%s handler never called", name) - } - } - - sopts := &ServerOptions{ - InitializedHandler: func(context.Context, *InitializedRequest) { - notificationChans["initialized"] <- 0 - }, - RootsListChangedHandler: func(context.Context, *RootsListChangedRequest) { - notificationChans["roots"] <- 0 - }, - ProgressNotificationHandler: func(context.Context, *ProgressNotificationServerRequest) { - notificationChans["progress_server"] <- 0 - }, - SubscribeHandler: func(context.Context, *SubscribeRequest) error { - notificationChans["subscribe"] <- 0 - return nil - }, - UnsubscribeHandler: func(context.Context, *UnsubscribeRequest) error { - notificationChans["unsubscribe"] <- 0 - return nil - }, - } - s := NewServer(testImpl, sopts) - AddTool(s, &Tool{ - Name: "greet", - Description: "say hi", - }, sayHi) - AddTool(s, &Tool{Name: "fail", InputSchema: &jsonschema.Schema{Type: "object"}}, - func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, any, error) { - return nil, nil, errTestFailure - }) - s.AddPrompt(codeReviewPrompt, codReviewPromptHandler) - s.AddPrompt(&Prompt{Name: "fail"}, func(_ context.Context, _ *GetPromptRequest) (*GetPromptResult, error) { - return nil, errTestFailure - }) - s.AddResource(resource1, readHandler) - s.AddResource(resource2, readHandler) - - // Connect the server. - ss, err := s.Connect(ctx, st, nil) - if err != nil { - t.Fatal(err) - } - if got := slices.Collect(s.Sessions()); len(got) != 1 { - t.Errorf("after connection, Clients() has length %d, want 1", len(got)) - } - - loggingMessages := make(chan *LoggingMessageParams, 100) // big enough for all logging - opts := &ClientOptions{ - CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { - return &CreateMessageResult{Model: "aModel", Content: &TextContent{}}, nil - }, - ElicitationHandler: func(ctx context.Context, req *ElicitRequest) (*ElicitResult, error) { - return &ElicitResult{ - Action: "accept", - Content: map[string]any{ - "name": "Test User", - "email": "test@example.com", - }, - }, nil - }, - ToolListChangedHandler: func(context.Context, *ToolListChangedRequest) { - notificationChans["tools"] <- 0 - }, - PromptListChangedHandler: func(context.Context, *PromptListChangedRequest) { - notificationChans["prompts"] <- 0 - }, - ResourceListChangedHandler: func(context.Context, *ResourceListChangedRequest) { - notificationChans["resources"] <- 0 - }, - LoggingMessageHandler: func(_ context.Context, req *LoggingMessageRequest) { - loggingMessages <- req.Params - }, - ProgressNotificationHandler: func(context.Context, *ProgressNotificationClientRequest) { - notificationChans["progress_client"] <- 0 - }, - ResourceUpdatedHandler: func(context.Context, *ResourceUpdatedNotificationRequest) { - notificationChans["resource_updated"] <- 0 - }, - } - c := NewClient(testImpl, opts) - rootAbs, err := filepath.Abs(filepath.FromSlash("testdata/files")) - if err != nil { - t.Fatal(err) - } - c.AddRoots(&Root{URI: "file://" + rootAbs}) - - // Connect the client. - cs, err := c.Connect(ctx, ct, nil) - if err != nil { - t.Fatal(err) - } - - waitForNotification(t, "initialized") - if err := cs.Ping(ctx, nil); err != nil { - t.Fatalf("ping failed: %v", err) - } - t.Run("prompts", func(t *testing.T) { - res, err := cs.ListPrompts(ctx, nil) - if err != nil { - t.Fatalf("prompts/list failed: %v", err) - } - wantPrompts := []*Prompt{ - { - Name: "code_review", - Description: "do a code review", - Arguments: []*PromptArgument{{Name: "Code", Required: true}}, - }, - {Name: "fail"}, - } - if diff := cmp.Diff(wantPrompts, res.Prompts); diff != "" { - t.Fatalf("prompts/list mismatch (-want +got):\n%s", diff) + // Channels to check if notification callbacks happened. + // These test asynchronous sending of notifications after a small delay (see + // Server.sendNotification). + notificationChans := map[string]chan int{} + for _, name := range []string{"initialized", "roots", "tools", "prompts", "resources", "progress_server", "progress_client", "resource_updated", "subscribe", "unsubscribe", "elicitation_complete"} { + notificationChans[name] = make(chan int, 1) } - gotReview, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "code_review", Arguments: map[string]string{"Code": "1+1"}}) - if err != nil { - t.Fatal(err) - } - wantReview := &GetPromptResult{ - Description: "Code review prompt", - Messages: []*PromptMessage{{ - Content: &TextContent{Text: "Please review the following code: 1+1"}, - Role: "user", - }}, - } - if diff := cmp.Diff(wantReview, gotReview); diff != "" { - t.Errorf("prompts/get 'code_review' mismatch (-want +got):\n%s", diff) - } - - if _, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "fail"}); err == nil || !strings.Contains(err.Error(), errTestFailure.Error()) { - t.Errorf("fail returned unexpected error: got %v, want containing %v", err, errTestFailure) - } - - s.AddPrompt(&Prompt{Name: "T"}, nil) - waitForNotification(t, "prompts") - s.RemovePrompts("T") - waitForNotification(t, "prompts") - }) - - t.Run("tools", func(t *testing.T) { - // ListTools is tested in client_list_test.go. - gotHi, err := cs.CallTool(ctx, &CallToolParams{ - Name: "greet", - Arguments: map[string]any{"Name": "user"}, - }) - if err != nil { - t.Fatal(err) - } - wantHi := &CallToolResult{ - Content: []Content{ - &TextContent{Text: "hi user"}, - }, - } - if diff := cmp.Diff(wantHi, gotHi, ctrCmpOpts...); diff != "" { - t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) - } - - gotFail, err := cs.CallTool(ctx, &CallToolParams{ - Name: "fail", - Arguments: map[string]any{}, - }) - // Counter-intuitively, when a tool fails, we don't expect an RPC error for - // call tool: instead, the failure is embedded in the result. - if err != nil { - t.Fatal(err) - } - wantFail := &CallToolResult{ - IsError: true, - Content: []Content{ - &TextContent{Text: errTestFailure.Error()}, - }, - } - if diff := cmp.Diff(wantFail, gotFail, ctrCmpOpts...); diff != "" { - t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff) - } - - // Check output schema validation. - badout := &Tool{ - Name: "badout", - OutputSchema: &jsonschema.Schema{ - Type: "object", - Properties: map[string]*jsonschema.Schema{ - "x": {Type: "string"}, - }, - }, - } - AddTool(s, badout, func(_ context.Context, _ *CallToolRequest, arg map[string]any) (*CallToolResult, map[string]any, error) { - return nil, map[string]any{"x": 1}, nil - }) - _, err = cs.CallTool(ctx, &CallToolParams{Name: "badout"}) - wantMsg := `has type "integer", want "string"` - if err == nil || !strings.Contains(err.Error(), wantMsg) { - t.Errorf("\ngot %q\nwant error message containing %q", err, wantMsg) - } - - // Check tools-changed notifications. - s.AddTool(&Tool{Name: "T", InputSchema: &jsonschema.Schema{Type: "object"}}, nopHandler) - waitForNotification(t, "tools") - s.RemoveTools("T") - waitForNotification(t, "tools") - }) - - t.Run("resources", func(t *testing.T) { - if runtime.GOOS == "windows" { - t.Skip("TODO: fix for Windows") - } - wantResources := []*Resource{resource2, resource1} - lrres, err := cs.ListResources(ctx, nil) - if err != nil { - t.Fatal(err) - } - if diff := cmp.Diff(wantResources, lrres.Resources); diff != "" { - t.Errorf("resources/list mismatch (-want, +got):\n%s", diff) - } - - template := &ResourceTemplate{ - Name: "rt", - MIMEType: "text/template", - URITemplate: "file:///{+filename}", // the '+' means that filename can contain '/' - } - s.AddResourceTemplate(template, readHandler) - tres, err := cs.ListResourceTemplates(ctx, nil) - if err != nil { - t.Fatal(err) - } - if diff := cmp.Diff([]*ResourceTemplate{template}, tres.ResourceTemplates); diff != "" { - t.Errorf("resources/list mismatch (-want, +got):\n%s", diff) - } - - for _, tt := range []struct { - uri string - mimeType string // "": not found; "text/plain": resource; "text/template": template - fail bool // non-nil error returned - }{ - {"file:///info.txt", "text/plain", false}, - {"file:///fail.txt", "", false}, - {"file:///template.txt", "text/template", false}, - {"file:///../private.txt", "", true}, // not found: escaping disallowed - } { - rres, err := cs.ReadResource(ctx, &ReadResourceParams{URI: tt.uri}) - if err != nil { - if code := errorCode(err); code == codeResourceNotFound { - if tt.mimeType != "" { - t.Errorf("%s: not found but expected it to be", tt.uri) - } - } else if !tt.fail { - t.Errorf("%s: unexpected error %v", tt.uri, err) - } - } else { - if tt.fail { - t.Errorf("%s: unexpected success", tt.uri) - } else if g, w := len(rres.Contents), 1; g != w { - t.Errorf("got %d contents, wanted %d", g, w) - } else { - c := rres.Contents[0] - if got := c.URI; got != tt.uri { - t.Errorf("got uri %q, want %q", got, tt.uri) - } - if got := c.MIMEType; got != tt.mimeType { - t.Errorf("%s: got MIME type %q, want %q", tt.uri, got, tt.mimeType) - } - } - } - } - - s.AddResource(&Resource{URI: "http://U"}, nil) - waitForNotification(t, "resources") - s.RemoveResources("http://U") - waitForNotification(t, "resources") - }) - t.Run("roots", func(t *testing.T) { - rootRes, err := ss.ListRoots(ctx, &ListRootsParams{}) - if err != nil { - t.Fatal(err) - } - gotRoots := rootRes.Roots - wantRoots := slices.Collect(c.roots.all()) - if diff := cmp.Diff(wantRoots, gotRoots); diff != "" { - t.Errorf("roots/list mismatch (-want +got):\n%s", diff) - } - - c.AddRoots(&Root{URI: "U"}) - waitForNotification(t, "roots") - c.RemoveRoots("U") - waitForNotification(t, "roots") - }) - t.Run("sampling", func(t *testing.T) { - // TODO: test that a client that doesn't have the handler returns CodeUnsupportedMethod. - res, err := ss.CreateMessage(ctx, &CreateMessageParams{}) - if err != nil { - t.Fatal(err) - } - if g, w := res.Model, "aModel"; g != w { - t.Errorf("got %q, want %q", g, w) - } - }) - t.Run("logging", func(t *testing.T) { - want := []*LoggingMessageParams{ - { - Logger: "test", - Level: "warning", - Data: map[string]any{ - "msg": "first", - "name": "Pat", - "logtest": true, - }, - }, - { - Logger: "test", - Level: "alert", - Data: map[string]any{ - "msg": "second", - "count": 2.0, - "logtest": true, - }, - }, - } - - check := func(t *testing.T) { + waitForNotification := func(t *testing.T, name string) { t.Helper() - var got []*LoggingMessageParams - // Read messages from this test until we've seen all we expect. - for len(got) < len(want) { - select { - case p := <-loggingMessages: - // Ignore logging from other tests. - if m, ok := p.Data.(map[string]any); ok && m["logtest"] != nil { - delete(m, "time") // remove time because it changes - got = append(got, p) - } - case <-time.After(time.Second): - t.Fatal("timed out waiting for log messages") - } - } - if diff := cmp.Diff(want, got); diff != "" { - t.Errorf("mismatch (-want, +got):\n%s", diff) - } + time.Sleep(notificationDelay * 2) + <-notificationChans[name] } - t.Run("direct", func(t *testing.T) { // Use the LoggingMessage method directly. + sopts := &ServerOptions{ + InitializedHandler: func(context.Context, *InitializedRequest) { + notificationChans["initialized"] <- 0 + }, + RootsListChangedHandler: func(context.Context, *RootsListChangedRequest) { + notificationChans["roots"] <- 0 + }, + ProgressNotificationHandler: func(context.Context, *ProgressNotificationServerRequest) { + notificationChans["progress_server"] <- 0 + }, + SubscribeHandler: func(context.Context, *SubscribeRequest) error { + notificationChans["subscribe"] <- 0 + return nil + }, + UnsubscribeHandler: func(context.Context, *UnsubscribeRequest) error { + notificationChans["unsubscribe"] <- 0 + return nil + }, + } + s := NewServer(testImpl, sopts) + AddTool(s, &Tool{ + Name: "greet", + Description: "say hi", + }, sayHi) + AddTool(s, &Tool{Name: "fail", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, any, error) { + return nil, nil, errTestFailure + }) + s.AddPrompt(codeReviewPrompt, codReviewPromptHandler) + s.AddPrompt(&Prompt{Name: "fail"}, func(_ context.Context, _ *GetPromptRequest) (*GetPromptResult, error) { + return nil, errTestFailure + }) + s.AddResource(resource1, readHandler) + s.AddResource(resource2, readHandler) - mustLog := func(level LoggingLevel, data any) { - t.Helper() - if err := ss.Log(ctx, &LoggingMessageParams{ - Logger: "test", - Level: level, - Data: data, - }); err != nil { - t.Fatal(err) - } + // Connect the server. + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + if got := slices.Collect(s.Sessions()); len(got) != 1 { + t.Errorf("after connection, Clients() has length %d, want 1", len(got)) + } + + loggingMessages := make(chan *LoggingMessageParams, 100) // big enough for all logging + opts := &ClientOptions{ + CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + return &CreateMessageResult{Model: "aModel", Content: &TextContent{}}, nil + }, + ElicitationHandler: func(ctx context.Context, req *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{ + Action: "accept", + Content: map[string]any{ + "name": "Test User", + "email": "test@example.com", + }, + }, nil + }, + ToolListChangedHandler: func(context.Context, *ToolListChangedRequest) { + notificationChans["tools"] <- 0 + }, + PromptListChangedHandler: func(context.Context, *PromptListChangedRequest) { + notificationChans["prompts"] <- 0 + }, + ResourceListChangedHandler: func(context.Context, *ResourceListChangedRequest) { + notificationChans["resources"] <- 0 + }, + LoggingMessageHandler: func(_ context.Context, req *LoggingMessageRequest) { + loggingMessages <- req.Params + }, + ProgressNotificationHandler: func(context.Context, *ProgressNotificationClientRequest) { + notificationChans["progress_client"] <- 0 + }, + ResourceUpdatedHandler: func(context.Context, *ResourceUpdatedNotificationRequest) { + notificationChans["resource_updated"] <- 0 + }, + ElicitationCompleteHandler: func(_ context.Context, req *ElicitationCompleteNotificationRequest) { + notificationChans["elicitation_complete"] <- 0 + }, + } + c := NewClient(testImpl, opts) + rootAbs, err := filepath.Abs(filepath.FromSlash("testdata/files")) + if err != nil { + t.Fatal(err) + } + c.AddRoots(&Root{URI: "file://" + rootAbs}) + + // Connect the client. + // + // Pin the session to 2025-11-25 so the legacy + // semantics apply. + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + + waitForNotification(t, "initialized") + if err := cs.Ping(ctx, nil); err != nil { + t.Fatalf("ping failed: %v", err) + } + + // ===== prompts ===== + t.Log("Testing prompts") + { + res, err := cs.ListPrompts(ctx, nil) + if err != nil { + t.Fatalf("prompts/list failed: %v", err) + } + wantPrompts := []*Prompt{ + { + Name: "code_review", + Description: "do a code review", + Arguments: []*PromptArgument{{Name: "Code", Required: true}}, + }, + {Name: "fail"}, + } + if diff := cmp.Diff(wantPrompts, res.Prompts); diff != "" { + t.Fatalf("prompts/list mismatch (-want +got):\n%s", diff) } - // Nothing should be logged until the client sets a level. - mustLog("info", "before") - if err := cs.SetLoggingLevel(ctx, &SetLoggingLevelParams{Level: "warning"}); err != nil { + gotReview, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "code_review", Arguments: map[string]string{"Code": "1+1"}}) + if err != nil { t.Fatal(err) } - mustLog("warning", want[0].Data) - mustLog("debug", "nope") // below the level - mustLog("info", "negative") // below the level - mustLog("alert", want[1].Data) - check(t) - }) + wantReview := &GetPromptResult{ + Description: "Code review prompt", + Messages: []*PromptMessage{{ + Content: &TextContent{Text: "Please review the following code: 1+1"}, + Role: "user", + }}, + } + if diff := cmp.Diff(wantReview, gotReview, ctrCmpOpts...); diff != "" { + t.Errorf("prompts/get 'code_review' mismatch (-want +got):\n%s", diff) + } - t.Run("handler", func(t *testing.T) { // Use the slog handler. - // We can't check the "before SetLevel" behavior because it's already been set. - // Not a big deal: that check is in LoggingMessage anyway. - logger := slog.New(NewLoggingHandler(ss, &LoggingHandlerOptions{LoggerName: "test"})) - logger.Warn("first", "name", "Pat", "logtest", true) - logger.Debug("nope") // below the level - logger.Info("negative") // below the level - logger.Log(ctx, LevelAlert, "second", "count", 2, "logtest", true) - check(t) - }) - }) - t.Run("progress", func(t *testing.T) { - ss.NotifyProgress(ctx, &ProgressNotificationParams{ - ProgressToken: "token-xyz", - Message: "progress update", - Progress: 0.5, - Total: 2, - }) - waitForNotification(t, "progress_client") - cs.NotifyProgress(ctx, &ProgressNotificationParams{ - ProgressToken: "token-abc", - Message: "progress update", - Progress: 1, - Total: 2, - }) - waitForNotification(t, "progress_server") - }) + if _, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "fail"}); err == nil || !strings.Contains(err.Error(), errTestFailure.Error()) { + t.Errorf("fail returned unexpected error: got %v, want containing %v", err, errTestFailure) + } - t.Run("resource_subscriptions", func(t *testing.T) { - err := cs.Subscribe(ctx, &SubscribeParams{ - URI: "test", - }) - if err != nil { - t.Fatal(err) + s.AddPrompt(&Prompt{Name: "T"}, nil) + waitForNotification(t, "prompts") + s.RemovePrompts("T") + waitForNotification(t, "prompts") } - waitForNotification(t, "subscribe") - s.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{ - URI: "test", - }) - waitForNotification(t, "resource_updated") - err = cs.Unsubscribe(ctx, &UnsubscribeParams{ - URI: "test", - }) - if err != nil { - t.Fatal(err) - } - waitForNotification(t, "unsubscribe") - // Verify the client does not receive the update after unsubscribing. - s.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{ - URI: "test", - }) - select { - case <-notificationChans["resource_updated"]: - t.Fatalf("resource updated after unsubscription") - case <-time.After(time.Second): + // ===== tools ===== + t.Log("Testing tools") + { + // ListTools is tested in client_list_test.go. + gotHi, err := cs.CallTool(ctx, &CallToolParams{ + Name: "greet", + Arguments: map[string]any{"Name": "user"}, + }) + if err != nil { + t.Fatal(err) + } + wantHi := &CallToolResult{ + Content: []Content{ + &TextContent{Text: "hi user"}, + }, + } + if diff := cmp.Diff(wantHi, gotHi, ctrCmpOpts...); diff != "" { + t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) + } + + gotFail, err := cs.CallTool(ctx, &CallToolParams{ + Name: "fail", + Arguments: map[string]any{}, + }) + // Counter-intuitively, when a tool fails, we don't expect an RPC error for + // call tool: instead, the failure is embedded in the result. + if err != nil { + t.Fatal(err) + } + wantFail := &CallToolResult{ + IsError: true, + Content: []Content{ + &TextContent{Text: errTestFailure.Error()}, + }, + } + if diff := cmp.Diff(wantFail, gotFail, ctrCmpOpts...); diff != "" { + t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff) + } + + // Check output schema validation. + badout := &Tool{ + Name: "badout", + OutputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "x": {Type: "string"}, + }, + }, + } + AddTool(s, badout, func(_ context.Context, _ *CallToolRequest, arg map[string]any) (*CallToolResult, map[string]any, error) { + return nil, map[string]any{"x": 1}, nil + }) + _, err = cs.CallTool(ctx, &CallToolParams{Name: "badout"}) + wantMsg := `has type "integer", want "string"` + if err == nil || !strings.Contains(err.Error(), wantMsg) { + t.Errorf("\ngot %q\nwant error message containing %q", err, wantMsg) + } + + // Check tools-changed notifications. + s.AddTool(&Tool{Name: "T", InputSchema: &jsonschema.Schema{Type: "object"}}, nopHandler) + waitForNotification(t, "tools") + s.RemoveTools("T") + waitForNotification(t, "tools") + } + + // ===== resources ===== + t.Log("Testing resources") + //TODO: fix for Windows + if runtime.GOOS != "windows" { + wantResources := []*Resource{resource2, resource1} + lrres, err := cs.ListResources(ctx, nil) + if err != nil { + t.Fatal(err) + } + if diff := cmp.Diff(wantResources, lrres.Resources); diff != "" { + t.Errorf("resources/list mismatch (-want, +got):\n%s", diff) + } + + template := &ResourceTemplate{ + Name: "rt", + MIMEType: "text/template", + URITemplate: "file:///{+filename}", // the '+' means that filename can contain '/' + } + s.AddResourceTemplate(template, readHandler) + tres, err := cs.ListResourceTemplates(ctx, nil) + if err != nil { + t.Fatal(err) + } + if diff := cmp.Diff([]*ResourceTemplate{template}, tres.ResourceTemplates); diff != "" { + t.Errorf("resources/list mismatch (-want, +got):\n%s", diff) + } + + for _, tt := range []struct { + uri string + mimeType string // "": not found; "text/plain": resource; "text/template": template + fail bool // non-nil error returned + }{ + {"file:///info.txt", "text/plain", false}, + {"file:///fail.txt", "", false}, + {"file:///template.txt", "text/template", false}, + {"file:///../private.txt", "", true}, // not found: escaping disallowed + } { + rres, err := cs.ReadResource(ctx, &ReadResourceParams{URI: tt.uri}) + if err != nil { + if code := errorCode(err); code == CodeResourceNotFound { + if tt.mimeType != "" { + t.Errorf("%s: not found but expected it to be", tt.uri) + } + } else if !tt.fail { + t.Errorf("%s: unexpected error %v", tt.uri, err) + } + } else { + if tt.fail { + t.Errorf("%s: unexpected success", tt.uri) + } else if g, w := len(rres.Contents), 1; g != w { + t.Errorf("got %d contents, wanted %d", g, w) + } else { + c := rres.Contents[0] + if got := c.URI; got != tt.uri { + t.Errorf("got uri %q, want %q", got, tt.uri) + } + if got := c.MIMEType; got != tt.mimeType { + t.Errorf("%s: got MIME type %q, want %q", tt.uri, got, tt.mimeType) + } + } + } + } + + s.AddResource(&Resource{URI: "http://U"}, nil) + waitForNotification(t, "resources") + s.RemoveResources("http://U") + waitForNotification(t, "resources") + } + + // ===== roots ===== + t.Log("Testing roots") + { + rootRes, err := ss.ListRoots(ctx, &ListRootsParams{}) + if err != nil { + t.Fatal(err) + } + gotRoots := rootRes.Roots + wantRoots := slices.Collect(c.roots.all()) + if diff := cmp.Diff(wantRoots, gotRoots); diff != "" { + t.Errorf("roots/list mismatch (-want +got):\n%s", diff) + } + + c.AddRoots(&Root{URI: "U"}) + waitForNotification(t, "roots") + c.RemoveRoots("U") + waitForNotification(t, "roots") + } + + // ===== sampling ===== + t.Log("Testing sampling") + { + // TODO: test that a client that doesn't have the handler returns CodeUnsupportedMethod. + res, err := ss.CreateMessage(ctx, &CreateMessageParams{}) + if err != nil { + t.Fatal(err) + } + if g, w := res.Model, "aModel"; g != w { + t.Errorf("got %q, want %q", g, w) + } + } + + // ===== logging ===== + t.Log("Testing logging") + { + want := []*LoggingMessageParams{ + { + Logger: "test", + Level: "warning", + Data: map[string]any{ + "msg": "first", + "name": "Pat", + "logtest": true, + }, + }, + { + Logger: "test", + Level: "alert", + Data: map[string]any{ + "msg": "second", + "count": 2.0, + "logtest": true, + }, + }, + } + + check := func(t *testing.T) { + t.Helper() + var got []*LoggingMessageParams + // Read messages from this test until we've seen all we expect. + for len(got) < len(want) { + p := <-loggingMessages + // Ignore logging from other tests. + if m, ok := p.Data.(map[string]any); ok && m["logtest"] != nil { + delete(m, "time") + got = append(got, p) + } + } + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("mismatch (-want, +got):\n%s", diff) + } + } + + // Use the LoggingMessage method directly. + t.Log("Testing logging (direct)") + { + mustLog := func(level LoggingLevel, data any) { + t.Helper() + if err := ss.Log(ctx, &LoggingMessageParams{ + Logger: "test", + Level: level, + Data: data, + }); err != nil { + t.Fatal(err) + } + } + + // Nothing should be logged until the client sets a level. + mustLog("info", "before") + if err := cs.SetLoggingLevel(ctx, &SetLoggingLevelParams{Level: "warning"}); err != nil { + t.Fatal(err) + } + mustLog("warning", want[0].Data) + mustLog("debug", "nope") // below the level + mustLog("info", "negative") // below the level + mustLog("alert", want[1].Data) + check(t) + } + + // Use the slog handler. + t.Log("Testing logging (handler)") + { + // We can't check the "before SetLevel" behavior because it's already been set. + // Not a big deal: that check is in LoggingMessage anyway. + logger := slog.New(NewLoggingHandler(ss, &LoggingHandlerOptions{LoggerName: "test"})) + logger.Warn("first", "name", "Pat", "logtest", true) + logger.Debug("nope") // below the level + logger.Info("negative") // below the level + logger.Log(ctx, LevelAlert, "second", "count", 2, "logtest", true) + check(t) + } + } + + // ===== progress ===== + t.Log("Testing progress") + { + ss.NotifyProgress(ctx, &ProgressNotificationParams{ + ProgressToken: "token-xyz", + Message: "progress update", + Progress: 0.5, + Total: 2, + }) + waitForNotification(t, "progress_client") + + cs.NotifyProgress(ctx, &ProgressNotificationParams{ + ProgressToken: "token-abc", + Message: "progress update", + Progress: 1, + Total: 2, + }) + waitForNotification(t, "progress_server") + } + + // ===== resource_subscriptions ===== + t.Log("Testing resource_subscriptions") + { + err := cs.Subscribe(ctx, &SubscribeParams{ + URI: "test", + }) + if err != nil { + t.Fatal(err) + } + synctest.Wait() + <-notificationChans["subscribe"] + + s.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{ + URI: "test", + }) + waitForNotification(t, "resource_updated") + + err = cs.Unsubscribe(ctx, &UnsubscribeParams{ + URI: "test", + }) + if err != nil { + t.Fatal(err) + } + waitForNotification(t, "unsubscribe") + + // Verify the client does not receive the update after unsubscribing. + s.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{ + URI: "test", + }) + synctest.Wait() + select { + case <-notificationChans["resource_updated"]: + t.Fatalf("resource updated after unsubscription") + default: + // Expected: no notification received + } + } + + // ===== elicitation ===== + t.Log("Testing elicitation") + { + result, err := ss.Elicit(ctx, &ElicitParams{ + Message: "Please provide information", + }) + if err != nil { + t.Fatal(err) + } + if result.Action != "accept" { + t.Errorf("got action %q, want %q", result.Action, "accept") + } + } + + // Disconnect. + cs.Close() + if err := ss.Wait(); err != nil { + t.Errorf("server failed: %v", err) + } + + // After disconnecting, neither client nor server should have any + // connections. + for range s.Sessions() { + t.Errorf("unexpected client after disconnection") } }) - - t.Run("elicitation", func(t *testing.T) { - result, err := ss.Elicit(ctx, &ElicitParams{ - Message: "Please provide information", - }) - if err != nil { - t.Fatal(err) - } - if result.Action != "accept" { - t.Errorf("got action %q, want %q", result.Action, "accept") - } - }) - - // Disconnect. - cs.Close() - if err := ss.Wait(); err != nil { - t.Errorf("server failed: %v", err) - } - - // After disconnecting, neither client nor server should have any - // connections. - for range s.Sessions() { - t.Errorf("unexpected client after disconnection") - } } // Registry of values to be referenced in tests. @@ -572,7 +614,7 @@ func errorCode(err error) int64 { if err == nil { return 0 } - var werr *jsonrpc2.WireError + var werr *jsonrpc.Error if errors.As(err, &werr) { return werr.Code } @@ -643,13 +685,11 @@ func TestServerClosing(t *testing.T) { ctx := context.Background() var wg sync.WaitGroup - wg.Add(1) - go func() { + wg.Go(func() { if err := cs.Wait(); err != nil { t.Errorf("server connection failed: %v", err) } - wg.Done() - }() + }) if _, err := cs.CallTool(ctx, &CallToolParams{ Name: "greet", Arguments: map[string]any{"Name": "user"}, @@ -666,73 +706,99 @@ func TestServerClosing(t *testing.T) { } } -func TestBatching(t *testing.T) { - ctx := context.Background() - ct, st := NewInMemoryTransports() - - s := NewServer(testImpl, nil) - _, err := s.Connect(ctx, st, nil) - if err != nil { - t.Fatal(err) - } - - c := NewClient(testImpl, nil) - // TODO: this test is broken, because increasing the batch size here causes - // 'initialize' to block. Therefore, we can only test with a size of 1. - // Since batching is being removed, we can probably just delete this. - const batchSize = 1 - cs, err := c.Connect(ctx, ct, nil) - if err != nil { - t.Fatal(err) - } - defer cs.Close() - - errs := make(chan error, batchSize) - for i := range batchSize { - go func() { - _, err := cs.ListTools(ctx, nil) - errs <- err - }() - time.Sleep(2 * time.Millisecond) - if i < batchSize-1 { - select { - case <-errs: - t.Errorf("ListTools: unexpected result for incomplete batch: %v", err) - default: - } - } - } -} - func TestCancellation(t *testing.T) { - var ( - start = make(chan struct{}) - cancelled = make(chan struct{}, 1) // don't block the request - ) - slowTool := func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { - start <- struct{}{} - select { - case <-ctx.Done(): - cancelled <- struct{}{} - case <-time.After(5 * time.Second): + synctest.Test(t, func(t *testing.T) { + var ( + start = make(chan struct{}) + cancelled = make(chan error, 1) // don't block the request + ) + slowTool := func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { + start <- struct{}{} + select { + case <-ctx.Done(): + cancelled <- context.Cause(ctx) + case <-time.After(5 * time.Second): + return nil, nil, nil + } return nil, nil, nil } - return nil, nil, nil - } - cs, _, cleanup := basicConnection(t, func(s *Server) { - AddTool(s, &Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, slowTool) - }) - defer cleanup() + cs, _, cleanup := basicConnection(t, func(s *Server) { + AddTool(s, &Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, slowTool) + }) + defer cleanup() - ctx, cancel := context.WithCancel(context.Background()) - go cs.CallTool(ctx, &CallToolParams{Name: "slow"}) - <-start - cancel() - select { - case <-cancelled: - case <-time.After(5 * time.Second): - t.Fatal("timeout waiting for cancellation") - } + ctx, cancel := context.WithCancel(context.Background()) + go cs.CallTool(ctx, &CallToolParams{Name: "slow"}) + <-start + cancel() + + // The client sends its context's error as the reason of its cancelled + // notification, and the handler reads it back as the cause. + cause := <-cancelled + if !errors.Is(cause, context.Canceled) { + t.Errorf("context.Cause = %v, want it to wrap context.Canceled", cause) + } + if want := "request cancelled by the peer: " + context.Canceled.Error(); cause == nil || cause.Error() != want { + t.Errorf("context.Cause = %v, want %q", cause, want) + } + }) +} + +// TestCancellationReason verifies that the reason a peer gives in its +// cancelled notification reaches the handler as the cause of its context, +// which is what lets a server log it as the specification asks. +// +// The call and the notification are written to the connection directly, so +// the reason is one the SDK's own client would never send; the session is +// pinned to 2025-11-25 so that a request written that way needs no _meta. +func TestCancellationReason(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + var ( + start = make(chan struct{}) + cancelled = make(chan error, 1) // don't block the request + ) + slowTool := func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { + start <- struct{}{} + select { + case <-ctx.Done(): + cancelled <- context.Cause(ctx) + case <-time.After(5 * time.Second): + cancelled <- nil + } + return nil, nil, nil + } + ctx := context.Background() + ct, st := NewInMemoryTransports() + s := NewServer(testImpl, nil) + AddTool(s, &Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, slowTool) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ss.Close() }) + cs, err := NewClient(testImpl, nil).Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = cs.Close() }) + + call := cs.conn.Call(ctx, methodCallTool, &CallToolParams{Name: "slow"}) + <-start + if err := cs.conn.Notify(ctx, notificationCancelled, &CancelledParams{RequestID: call.ID().Raw(), Reason: "user asked"}); err != nil { + t.Fatal(err) + } + + cause := <-cancelled + if cause == nil { + t.Fatal("the tool ran to completion, want it cancelled") + } + if !errors.Is(cause, context.Canceled) { + t.Errorf("context.Cause = %v, want it to wrap context.Canceled", cause) + } + if got, want := cause.Error(), "request cancelled by the peer: user asked"; got != want { + t.Errorf("context.Cause = %q, want %q", got, want) + } + }) } func TestMiddleware(t *testing.T) { @@ -760,7 +826,10 @@ func TestMiddleware(t *testing.T) { c.AddSendingMiddleware(traceCalls[*ClientSession](&cbuf, "S1"), traceCalls[*ClientSession](&cbuf, "S2")) c.AddReceivingMiddleware(traceCalls[*ClientSession](&cbuf, "R1"), traceCalls[*ClientSession](&cbuf, "R2")) - cs, err := c.Connect(ctx, ct, nil) + // Pin to 2025-11-25 because the test's expected wire sequence asserts + // the legacy initialize / notifications/initialized handshake, which + // 2026-07-28 replaces with server/discover. + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatal(err) } @@ -850,7 +919,7 @@ func TestNoJSONNull(t *testing.T) { } c := NewClient(testImpl, nil) - cs, err := c.Connect(ctx, ct, nil) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatal(err) } @@ -898,53 +967,107 @@ func nopHandler(context.Context, *CallToolRequest) (*CallToolResult, error) { } func TestKeepAlive(t *testing.T) { - // TODO: try to use the new synctest package for this test once we upgrade to Go 1.24+. - // synctest would allow us to control time and avoid the time.Sleep calls, making the test - // faster and more deterministic. - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() - ct, st := NewInMemoryTransports() + ct, st := NewInMemoryTransports() - serverOpts := &ServerOptions{ - KeepAlive: 100 * time.Millisecond, - } - s := NewServer(testImpl, serverOpts) - AddTool(s, greetTool(), sayHi) + serverOpts := &ServerOptions{ + KeepAlive: 100 * time.Millisecond, + } + s := NewServer(testImpl, serverOpts) + AddTool(s, greetTool(), sayHi) - ss, err := s.Connect(ctx, st, nil) - if err != nil { - t.Fatal(err) - } - defer ss.Close() + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() - clientOpts := &ClientOptions{ - KeepAlive: 100 * time.Millisecond, - } - c := NewClient(testImpl, clientOpts) - cs, err := c.Connect(ctx, ct, nil) - if err != nil { - t.Fatal(err) - } - defer cs.Close() + clientOpts := &ClientOptions{ + KeepAlive: 100 * time.Millisecond, + } + c := NewClient(testImpl, clientOpts) + cs, err := c.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() - // Wait for a few keepalive cycles to ensure pings are working - time.Sleep(300 * time.Millisecond) + // Wait for a few keepalive cycles to ensure pings are working. + // With synctest, this advances simulated time instantly. + time.Sleep(300 * time.Millisecond) - // Test that the connection is still alive by making a call - result, err := cs.CallTool(ctx, &CallToolParams{ - Name: "greet", - Arguments: map[string]any{"Name": "user"}, + // Test that the connection is still alive by making a call + result, err := cs.CallTool(ctx, &CallToolParams{ + Name: "greet", + Arguments: map[string]any{"Name": "user"}, + }) + if err != nil { + t.Fatalf("call failed after keepalive: %v", err) + } + if len(result.Content) == 0 { + t.Fatal("expected content in result") + } + if textContent, ok := result.Content[0].(*TextContent); !ok || textContent.Text != "hi user" { + t.Fatalf("unexpected result: %v", result.Content[0]) + } + }) +} + +func TestKeepAliveMethodNotFound(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + + ct, st := NewInMemoryTransports() + + // Server that rejects ping with method-not-found, simulating a + // server that does not implement the optional ping method. + s := NewServer(testImpl, nil) + AddTool(s, greetTool(), sayHi) + s.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == "ping" { + return nil, jsonrpc2.ErrMethodNotFound + } + return next(ctx, method, req) + } + }) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + clientOpts := &ClientOptions{ + KeepAlive: 50 * time.Millisecond, + } + c := NewClient(testImpl, clientOpts) + cs, err := c.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Advance past several keepalive cycles. + time.Sleep(200 * time.Millisecond) + + // The session should still be alive despite the server not + // supporting ping. + result, err := cs.CallTool(ctx, &CallToolParams{ + Name: "greet", + Arguments: map[string]any{"Name": "user"}, + }) + if err != nil { + t.Fatalf("call failed after keepalive with method-not-found: %v", err) + } + if len(result.Content) == 0 { + t.Fatal("expected content in result") + } + if textContent, ok := result.Content[0].(*TextContent); !ok || textContent.Text != "hi user" { + t.Fatalf("unexpected result: %v", result.Content[0]) + } }) - if err != nil { - t.Fatalf("call failed after keepalive: %v", err) - } - if len(result.Content) == 0 { - t.Fatal("expected content in result") - } - if textContent, ok := result.Content[0].(*TextContent); !ok || textContent.Text != "hi user" { - t.Fatalf("unexpected result: %v", result.Content[0]) - } } func TestElicitationUnsupportedMethod(t *testing.T) { @@ -965,7 +1088,7 @@ func TestElicitationUnsupportedMethod(t *testing.T) { return &CreateMessageResult{Model: "aModel", Content: &TextContent{}}, nil }, }) - cs, err := c.Connect(ctx, ct, nil) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatal(err) } @@ -985,14 +1108,19 @@ func TestElicitationUnsupportedMethod(t *testing.T) { if err == nil { t.Error("expected error when ElicitationHandler is not provided, got nil") } - if code := errorCode(err); code != codeUnsupportedMethod { - t.Errorf("got error code %d, want %d (CodeUnsupportedMethod)", code, codeUnsupportedMethod) + if code := errorCode(err); code != -1 { + t.Errorf("got error code %d, want -1", code) } if !strings.Contains(err.Error(), "does not support elicitation") { t.Errorf("error should mention unsupported elicitation, got: %v", err) } } +func anyPtr[T any](v T) *any { + var a any = v + return &a +} + func TestElicitationSchemaValidation(t *testing.T) { ctx := context.Background() ct, st := NewInMemoryTransports() @@ -1009,7 +1137,7 @@ func TestElicitationSchemaValidation(t *testing.T) { return &ElicitResult{Action: "accept", Content: map[string]any{"test": "value"}}, nil }, }) - cs, err := c.Connect(ctx, ct, nil) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatal(err) } @@ -1080,7 +1208,7 @@ func TestElicitationSchemaValidation(t *testing.T) { }, }, { - name: "string enum", + name: "single select untitled enum", schema: &jsonschema.Schema{ Type: "object", Properties: map[string]*jsonschema.Schema{ @@ -1096,7 +1224,7 @@ func TestElicitationSchemaValidation(t *testing.T) { }, }, { - name: "enum with matching enumNames", + name: "legacy titled enum", schema: &jsonschema.Schema{ Type: "object", Properties: map[string]*jsonschema.Schema{ @@ -1114,6 +1242,77 @@ func TestElicitationSchemaValidation(t *testing.T) { }, }, }, + { + name: "single select titled enum", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "priority": { + Type: "string", + OneOf: []*jsonschema.Schema{ + { + Const: anyPtr("high"), + Title: "High Priority", + }, + { + Const: anyPtr("medium"), + Title: "Medium Priority", + }, + { + Const: anyPtr("low"), + Title: "Low Priority", + }, + }, + }, + }, + }, + }, + { + name: "multi select untitled enum", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "status": { + Type: "array", + Items: &jsonschema.Schema{ + Type: "string", + Enum: []any{ + "active", + "inactive", + "pending", + }, + }, + }, + }, + }, + }, + { + name: "multi select titled enum", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "priority": { + Type: "array", + Items: &jsonschema.Schema{ + AnyOf: []*jsonschema.Schema{ + { + Const: anyPtr("high"), + Title: "High Priority", + }, + { + Const: anyPtr("medium"), + Title: "Medium Priority", + }, + { + Const: anyPtr("low"), + Title: "Low Priority", + }, + }, + }, + }, + }, + }, + }, } for _, tc := range validSchemas { @@ -1164,17 +1363,7 @@ func TestElicitationSchemaValidation(t *testing.T) { "config": {Type: "object"}, }, }, - expectedError: "elicit schema property \"config\" has unsupported type \"object\", only string, number, integer, and boolean are allowed", - }, - { - name: "array property", - schema: &jsonschema.Schema{ - Type: "object", - Properties: map[string]*jsonschema.Schema{ - "tags": {Type: "array", Items: &jsonschema.Schema{Type: "string"}}, - }, - }, - expectedError: "elicit schema property \"tags\" has unsupported type \"array\", only string, number, integer, and boolean are allowed", + expectedError: "elicit schema property \"config\" has unsupported type \"object\", only string, number, integer, boolean, and array are allowed", }, { name: "array without items", @@ -1184,7 +1373,27 @@ func TestElicitationSchemaValidation(t *testing.T) { "items": {Type: "array"}, }, }, - expectedError: "elicit schema property \"items\" has unsupported type \"array\", only string, number, integer, and boolean are allowed", + expectedError: "elicit schema property \"items\" is array but missing 'items' definition", + }, + { + name: "array with integer items", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "items": {Type: "array", Items: &jsonschema.Schema{Type: "integer"}}, + }, + }, + expectedError: "elicit schema property \"items\" items have unsupported type \"integer\"", + }, + { + name: "array of generic strings", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "items": {Type: "array", Items: &jsonschema.Schema{Type: "string"}}, + }, + }, + expectedError: "elicit schema property \"items\" items must specify enum for untitled enums", }, { name: "unsupported string format", @@ -1254,7 +1463,37 @@ func TestElicitationSchemaValidation(t *testing.T) { "enabled": {Type: "boolean", Default: json.RawMessage(`"not-a-boolean"`)}, }, }, - expectedError: "elicit schema property \"enabled\" has invalid default value, must be a boolean", + expectedError: "elicit schema property \"enabled\" has invalid default value, must be a bool", + }, + { + name: "string with invalid default", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "enabled": {Type: "string", Default: json.RawMessage("true")}, + }, + }, + expectedError: "elicit schema property \"enabled\" has invalid default value, must be a string", + }, + { + name: "integer with invalid default", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "enabled": {Type: "integer", Default: json.RawMessage("true")}, + }, + }, + expectedError: "elicit schema property \"enabled\" has default value that cannot be interpreted as an int or float", + }, + { + name: "number with invalid default", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "enabled": {Type: "number", Default: json.RawMessage("true")}, + }, + }, + expectedError: "elicit schema property \"enabled\" has default value that cannot be interpreted as an int or float", }, { name: "enum with mismatched enumNames length", @@ -1309,7 +1548,29 @@ func TestElicitationSchemaValidation(t *testing.T) { }, }, }, - expectedError: "elicit schema property \"priority\" has unsupported type \"\", only string, number, integer, and boolean are allowed", + expectedError: "elicit schema property \"priority\" has unsupported type \"\", only string, number, integer, boolean, and array are allowed", + }, + { + name: "titled enum without const", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "priority": { + Type: "string", + OneOf: []*jsonschema.Schema{ + { + Const: anyPtr("high"), + Title: "High Priority", + }, + { + Type: "string", + Title: "Other Priority", + }, + }, + }, + }, + }, + expectedError: "elicit schema property \"priority\" oneOf has invalid entry: const is required for titled enum entries", }, { name: "untyped property", @@ -1319,7 +1580,7 @@ func TestElicitationSchemaValidation(t *testing.T) { "data": {}, }, }, - expectedError: "elicit schema property \"data\" has unsupported type \"\", only string, number, integer, and boolean are allowed", + expectedError: "elicit schema property \"data\" has unsupported type \"\", only string, number, integer, boolean, and array are allowed", }, } @@ -1333,8 +1594,8 @@ func TestElicitationSchemaValidation(t *testing.T) { t.Errorf("expected error for invalid schema %q, got nil", tc.name) return } - if code := errorCode(err); code != codeInvalidParams { - t.Errorf("got error code %d, want %d (CodeInvalidParams)", code, codeInvalidParams) + if code := errorCode(err); code != jsonrpc.CodeInvalidParams { + t.Errorf("got error code %d, want %d (CodeInvalidParams)", code, jsonrpc.CodeInvalidParams) } if !strings.Contains(err.Error(), tc.expectedError) { t.Errorf("error message %q does not contain expected text %q", err.Error(), tc.expectedError) @@ -1343,6 +1604,92 @@ func TestElicitationSchemaValidation(t *testing.T) { } } +func TestElicitContentValidation(t *testing.T) { + ctx := context.Background() + ct, st := NewInMemoryTransports() + + s := NewServer(testImpl, nil) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + // Set up a client that exercises valid/invalid elicitation: the returned + // Content from the handler ("potato") is validated against the schemas + // defined in the testcases below. + c := NewClient(testImpl, &ClientOptions{ + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept", Content: map[string]any{"test": "potato"}}, nil + }, + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + testcases := []struct { + name string + schema *jsonschema.Schema + expectedError string + }{ + { + name: "string enum with schema not matching content", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "test": { + Type: "string", + OneOf: []*jsonschema.Schema{ + { + Const: anyPtr("high"), + Title: "High Priority", + }, + }, + }, + }, + }, + expectedError: "oneOf: did not validate against any of", + }, + { + name: "string enum with schema matching content", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "test": { + Type: "string", + OneOf: []*jsonschema.Schema{ + { + Const: anyPtr("potato"), + Title: "Potato Priority", + }, + }, + }, + }, + }, + expectedError: "", + }, + } + for _, tc := range testcases { + t.Run(tc.name, func(t *testing.T) { + _, err := ss.Elicit(ctx, &ElicitParams{ + Message: "Test schema: " + tc.name, + RequestedSchema: tc.schema, + }) + if tc.expectedError != "" { + if err == nil { + t.Errorf("expected error but got no error: %s", tc.expectedError) + return + } + if !strings.Contains(err.Error(), tc.expectedError) { + t.Errorf("error message %q does not contain expected text %q", err.Error(), tc.expectedError) + } + } + }) + } +} + func TestElicitationProgressToken(t *testing.T) { ctx := context.Background() ct, st := NewInMemoryTransports() @@ -1359,7 +1706,7 @@ func TestElicitationProgressToken(t *testing.T) { return &ElicitResult{Action: "accept"}, nil }, }) - cs, err := c.Connect(ctx, ct, nil) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatal(err) } @@ -1384,7 +1731,7 @@ func TestElicitationProgressToken(t *testing.T) { func TestElicitationCapabilityDeclaration(t *testing.T) { ctx := context.Background() - t.Run("with_handler", func(t *testing.T) { + t.Run("with handler", func(t *testing.T) { ct, st := NewInMemoryTransports() // Client with ElicitationHandler should declare capability @@ -1401,7 +1748,7 @@ func TestElicitationCapabilityDeclaration(t *testing.T) { } defer ss.Close() - cs, err := c.Connect(ctx, ct, nil) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatal(err) } @@ -1421,7 +1768,7 @@ func TestElicitationCapabilityDeclaration(t *testing.T) { } }) - t.Run("without_handler", func(t *testing.T) { + t.Run("without handler", func(t *testing.T) { ct, st := NewInMemoryTransports() // Client without ElicitationHandler should not declare capability @@ -1438,7 +1785,7 @@ func TestElicitationCapabilityDeclaration(t *testing.T) { } defer ss.Close() - cs, err := c.Connect(ctx, ct, nil) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatal(err) } @@ -1453,58 +1800,278 @@ func TestElicitationCapabilityDeclaration(t *testing.T) { if err == nil { t.Error("expected UnsupportedMethod error when no capability declared") } - if code := errorCode(err); code != codeUnsupportedMethod { - t.Errorf("got error code %d, want %d (CodeUnsupportedMethod)", code, codeUnsupportedMethod) + if code := errorCode(err); code != -1 { + t.Errorf("got error code %d, want -1", code) } }) } -func TestKeepAliveFailure(t *testing.T) { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - +func TestElicitationDefaultValues(t *testing.T) { + ctx := context.Background() ct, st := NewInMemoryTransports() - // Server without keepalive (to test one-sided keepalive) s := NewServer(testImpl, nil) - AddTool(s, greetTool(), sayHi) ss, err := s.Connect(ctx, st, nil) if err != nil { t.Fatal(err) } + defer ss.Close() - // Client with short keepalive - clientOpts := &ClientOptions{ - KeepAlive: 50 * time.Millisecond, - } - c := NewClient(testImpl, clientOpts) - cs, err := c.Connect(ctx, ct, nil) + c := NewClient(testImpl, &ClientOptions{ + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept", Content: map[string]any{"default": "response"}}, nil + }, + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatal(err) } defer cs.Close() - // Let the connection establish properly first - time.Sleep(30 * time.Millisecond) + testcases := []struct { + name string + schema *jsonschema.Schema + expected map[string]any + }{ + { + name: "boolean with default", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "key": {Type: "boolean", Default: json.RawMessage("true")}, + }, + }, + expected: map[string]any{"key": true, "default": "response"}, + }, + { + name: "string with default", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "key": {Type: "string", Default: json.RawMessage("\"potato\"")}, + }, + }, + expected: map[string]any{"key": "potato", "default": "response"}, + }, + { + name: "integer with default", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "key": {Type: "integer", Default: json.RawMessage("123")}, + }, + }, + expected: map[string]any{"key": float64(123), "default": "response"}, + }, + { + name: "number with default", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "key": {Type: "number", Default: json.RawMessage("89.7")}, + }, + }, + expected: map[string]any{"key": float64(89.7), "default": "response"}, + }, + { + name: "enum with default", + schema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "key": {Type: "string", Enum: []any{"one", "two"}, Default: json.RawMessage("\"one\"")}, + }, + }, + expected: map[string]any{"key": "one", "default": "response"}, + }, + } + for _, tc := range testcases { + t.Run(tc.name, func(t *testing.T) { + res, err := ss.Elicit(ctx, &ElicitParams{ + Message: "Test schema with defaults: " + tc.name, + RequestedSchema: tc.schema, + }) + if err != nil { + t.Fatalf("expected no error for default schema %q, got: %v", tc.name, err) + } + if diff := cmp.Diff(tc.expected, res.Content); diff != "" { + t.Errorf("%s: did not get expected value, -want +got:\n%s", tc.name, diff) + } + }) + } +} - // simulate ping failure - ss.Close() +func TestKeepAliveFailure(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + + // This test doesn't pinpoint keepalive detected failures well due to the fact that in-memory transports + // propagate `io.EOF` synchronously, causing the transport level connection to be closed immediately. + // TODO: propose better transports that would allow testing this scenario precisely. + ct, st := NewInMemoryTransports() + + // Server without keepalive (to test one-sided keepalive) + s := NewServer(testImpl, nil) + AddTool(s, greetTool(), sayHi) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + + // Client with short keepalive + clientOpts := &ClientOptions{ + KeepAlive: 50 * time.Millisecond, + } + c := NewClient(testImpl, clientOpts) + cs, err := c.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Let the connection establish properly first + synctest.Wait() + + // simulate ping failure + ss.Close() + + time.Sleep(100 * time.Millisecond) + synctest.Wait() - // Wait for keepalive to detect the failure and close the client - // check periodically instead of just waiting - deadline := time.Now().Add(1 * time.Second) - for time.Now().Before(deadline) { _, err = cs.CallTool(ctx, &CallToolParams{ Name: "greet", Arguments: map[string]any{"Name": "user"}, }) - if errors.Is(err, ErrConnectionClosed) { + if err != nil && (errors.Is(err, ErrConnectionClosed) || strings.Contains(err.Error(), "connection closed")) { return // Test passed } - time.Sleep(25 * time.Millisecond) - } - t.Errorf("expected connection to be closed by keepalive, but it wasn't. Last error: %v", err) + t.Errorf("expected connection to be closed by keepalive, but it wasn't. Last error: %v", err) + }) +} + +// TestKeepAliveFailure_Logged verifies that a keepalive ping failure is +// reported via the configured slog.Logger instead of being silently dropped. +// Regression test for #218. +func TestKeepAliveFailure_Logged(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + + ct, st := NewInMemoryTransports() + + // Server without keepalive. + s := NewServer(testImpl, nil) + AddTool(s, greetTool(), sayHi) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + + // Client with a short keepalive and a logger that writes to a + // buffer so we can assert on its output. + var buf bytes.Buffer + clientOpts := &ClientOptions{ + KeepAlive: 50 * time.Millisecond, + Logger: slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelError})), + } + c := NewClient(testImpl, clientOpts) + // Pin to 2025-11-25: KeepAlive uses the ping RPC, which is removed + // in 2026-07-28, so keepalive is only meaningful on legacy protocol + // versions. + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + synctest.Wait() + + // Trigger a ping failure by closing the server side. + ss.Close() + + time.Sleep(100 * time.Millisecond) + synctest.Wait() + + got := buf.String() // slog serializes Write calls internally + if !strings.Contains(got, "keepalive ping failed") { + t.Errorf("expected keepalive failure to be logged, got log output:\n%s", got) + } + }) +} + +// scriptedKeepaliveSession is a keepaliveSession test double whose Ping +// returns errors from a script (one entry consumed per call; the last entry +// repeats once exhausted), and records how many times Close was called. Ping +// returns immediately so the keepalive loop's pace is driven purely by the +// ticker, making the test deterministic under synctest. +type scriptedKeepaliveSession struct { + pingErrs []error + pingCalls atomic.Int64 + closeCalls atomic.Int64 +} + +func (s *scriptedKeepaliveSession) Ping(context.Context, *PingParams) error { + n := int(s.pingCalls.Add(1)) - 1 + if n >= len(s.pingErrs) { + n = len(s.pingErrs) - 1 + } + return s.pingErrs[n] +} + +func (s *scriptedKeepaliveSession) Close() error { + s.closeCalls.Add(1) + return nil +} + +// TestStartKeepalive_FailureThreshold verifies that the session is kept alive +// across consecutive ping failures below the threshold and only closed once the +// threshold is reached. +func TestStartKeepalive_FailureThreshold(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + const interval = 100 * time.Millisecond + sess := &scriptedKeepaliveSession{pingErrs: []error{errors.New("boom")}} + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + var cancel context.CancelFunc + startKeepalive(sess, interval, 3, &cancel, logger) + defer cancel() + + // After two ticks → two failures, still below threshold 3: not closed. + time.Sleep(2*interval + interval/2) + synctest.Wait() + if got := sess.closeCalls.Load(); got != 0 { + t.Fatalf("session closed below threshold: closeCalls=%d (pingCalls=%d)", got, sess.pingCalls.Load()) + } + + // Third tick → third failure reaches threshold: session closed. + time.Sleep(interval) + synctest.Wait() + if got := sess.closeCalls.Load(); got != 1 { + t.Fatalf("expected one Close at threshold, got closeCalls=%d (pingCalls=%d)", got, sess.pingCalls.Load()) + } + }) +} + +// TestStartKeepalive_SuccessResetsFailures verifies that a successful ping +// resets the consecutive-failure counter, so an isolated failure between +// successes never accumulates toward the threshold. +func TestStartKeepalive_SuccessResetsFailures(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + const interval = 100 * time.Millisecond + // fail, success, fail, fail, then success (the tail repeats): the run + // never has 3 consecutive failures, so the session is never closed. + sess := &scriptedKeepaliveSession{pingErrs: []error{ + errors.New("boom"), nil, errors.New("boom"), errors.New("boom"), nil, + }} + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + var cancel context.CancelFunc + startKeepalive(sess, interval, 3, &cancel, logger) + defer cancel() + + time.Sleep(6 * interval) + synctest.Wait() + if got := sess.closeCalls.Load(); got != 0 { + t.Fatalf("session closed despite a success resetting the counter: closeCalls=%d (pingCalls=%d)", got, sess.pingCalls.Load()) + } + }) } func TestAddTool_DuplicateNoPanicAndNoDuplicate(t *testing.T) { @@ -1543,98 +2110,125 @@ func TestAddTool_DuplicateNoPanicAndNoDuplicate(t *testing.T) { } func TestSynchronousNotifications(t *testing.T) { - var toolsChanged atomic.Bool - clientOpts := &ClientOptions{ - ToolListChangedHandler: func(ctx context.Context, req *ToolListChangedRequest) { - toolsChanged.Store(true) - }, - CreateMessageHandler: func(ctx context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) { - if !toolsChanged.Load() { - return nil, fmt.Errorf("didn't get a tools changed notification") - } - // TODO(rfindley): investigate the error returned from this test if - // CreateMessageResult is new(CreateMessageResult): it's a mysterious - // unmarshalling error that we should improve. - return &CreateMessageResult{Content: &TextContent{}}, nil - }, - } - client := NewClient(testImpl, clientOpts) + synctest.Test(t, func(t *testing.T) { + var toolsChanged atomic.Int32 + clientOpts := &ClientOptions{ + ToolListChangedHandler: func(ctx context.Context, req *ToolListChangedRequest) { + toolsChanged.Add(1) + }, + CreateMessageHandler: func(ctx context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) { + // See the comment after "from server" below. + if n := toolsChanged.Load(); n != 1 { + return nil, fmt.Errorf("got %d tools-changed notification, wanted 1", n) + } + // TODO(rfindley): investigate the error returned from this test if + // CreateMessageResult is new(CreateMessageResult): it's a mysterious + // unmarshalling error that we should improve. + return &CreateMessageResult{Content: &TextContent{}}, nil + }, + } + client := NewClient(testImpl, clientOpts) - var rootsChanged atomic.Bool - serverOpts := &ServerOptions{ - RootsListChangedHandler: func(_ context.Context, req *RootsListChangedRequest) { - rootsChanged.Store(true) - }, - } - server := NewServer(testImpl, serverOpts) - cs, ss, cleanup := basicClientServerConnection(t, client, server, func(s *Server) { - AddTool(s, &Tool{Name: "tool"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { - if !rootsChanged.Load() { - return nil, nil, fmt.Errorf("didn't get root change notification") - } - return new(CallToolResult), nil, nil - }) - }) - defer cleanup() - - t.Run("from client", func(t *testing.T) { - client.AddRoots(&Root{Name: "myroot", URI: "file://foo"}) - res, err := cs.CallTool(context.Background(), &CallToolParams{Name: "tool"}) + var rootsChanged atomic.Bool + serverOpts := &ServerOptions{ + RootsListChangedHandler: func(_ context.Context, req *RootsListChangedRequest) { + rootsChanged.Store(true) + }, + } + server := NewServer(testImpl, serverOpts) + addTool := func(s *Server) { + AddTool(s, &Tool{Name: "tool"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { + if !rootsChanged.Load() { + return nil, nil, fmt.Errorf("didn't get root change notification") + } + return new(CallToolResult), nil, nil + }) + } + ctx := context.Background() + ct, st := NewInMemoryTransports() + addTool(server) + ss, err := server.Connect(ctx, st, nil) if err != nil { - t.Fatalf("CallTool failed: %v", err) + t.Fatal(err) } - if res.IsError { - t.Errorf("tool error: %v", res.Content[0].(*TextContent).Text) + t.Cleanup(func() { _ = ss.Close() }) + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) } - }) + t.Cleanup(func() { + _ = cs.Close() + ss.Wait() + }) - t.Run("from server", func(t *testing.T) { - server.RemoveTools("tool") - if _, err := ss.CreateMessage(context.Background(), new(CreateMessageParams)); err != nil { - t.Errorf("CreateMessage failed: %v", err) + t.Log("from client") + { + client.AddRoots(&Root{Name: "myroot", URI: "file://foo"}) + res, err := cs.CallTool(context.Background(), &CallToolParams{Name: "tool"}) + if err != nil { + t.Fatalf("CallTool failed: %v", err) + } + if res.IsError { + t.Errorf("tool error: %v", res.Content[0].(*TextContent).Text) + } } + + t.Log("from server") + { + // Despite all this tool-changed activity, we expect only one notification. + for range 10 { + server.RemoveTools("tool") + addTool(server) + } + + time.Sleep(notificationDelay * 2) // Wait for delayed notification. + if _, err := ss.CreateMessage(context.Background(), new(CreateMessageParams)); err != nil { + t.Errorf("CreateMessage failed: %v", err) + } + } + }) } func TestNoDistributedDeadlock(t *testing.T) { - // This test verifies that calls are asynchronous, and so it's not possible - // to have a distributed deadlock. - // - // The setup creates potential deadlock for both the client and server: the - // client sends a call to tool1, which itself calls createMessage, which in - // turn calls tool2, which calls ping. - // - // If the server were not asynchronous, the call to tool2 would hang. If the - // client were not asynchronous, the call to ping would hang. - // - // Such a scenario is unlikely in practice, but is still theoretically - // possible, and in any case making tool calls asynchronous by default - // delegates synchronization to the user. - clientOpts := &ClientOptions{ - CreateMessageHandler: func(ctx context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) { - req.Session.CallTool(ctx, &CallToolParams{Name: "tool2"}) - return &CreateMessageResult{Content: &TextContent{}}, nil - }, - } - client := NewClient(testImpl, clientOpts) - cs, _, cleanup := basicClientServerConnection(t, client, nil, func(s *Server) { - AddTool(s, &Tool{Name: "tool1"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { - req.Session.CreateMessage(ctx, new(CreateMessageParams)) - return new(CallToolResult), nil, nil + synctest.Test(t, func(t *testing.T) { + // This test verifies that calls are asynchronous, and so it's not possible + // to have a distributed deadlock. + // + // The setup creates potential deadlock for both the client and server: the + // client sends a call to tool1, which itself calls createMessage, which in + // turn calls tool2, which calls ping. + // + // If the server were not asynchronous, the call to tool2 would hang. If the + // client were not asynchronous, the call to ping would hang. + // + // Such a scenario is unlikely in practice, but is still theoretically + // possible, and in any case making tool calls asynchronous by default + // delegates synchronization to the user. + clientOpts := &ClientOptions{ + CreateMessageHandler: func(ctx context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) { + req.Session.CallTool(ctx, &CallToolParams{Name: "tool2"}) + return &CreateMessageResult{Content: &TextContent{}}, nil + }, + } + client := NewClient(testImpl, clientOpts) + cs, _, cleanup := basicClientServerConnection(t, client, nil, func(s *Server) { + AddTool(s, &Tool{Name: "tool1"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { + req.Session.CreateMessage(ctx, new(CreateMessageParams)) + return new(CallToolResult), nil, nil + }) + AddTool(s, &Tool{Name: "tool2"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { + req.Session.Ping(ctx, nil) + return new(CallToolResult), nil, nil + }) }) - AddTool(s, &Tool{Name: "tool2"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { - req.Session.Ping(ctx, nil) - return new(CallToolResult), nil, nil - }) - }) - defer cleanup() + defer cleanup() - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - if _, err := cs.CallTool(ctx, &CallToolParams{Name: "tool1"}); err != nil { - // should not deadlock - t.Fatalf("CallTool failed: %v", err) - } + if _, err := cs.CallTool(context.Background(), &CallToolParams{Name: "tool1"}); err != nil { + // should not deadlock + t.Fatalf("CallTool failed: %v", err) + } + }) } var testImpl = &Implementation{Name: "test", Version: "v1.0.0"} @@ -1860,7 +2454,7 @@ func TestToolErrorMiddleware(t *testing.T) { res, err := h(ctx, method, req) if err == nil { if ctr, ok := res.(*CallToolResult); ok { - middleErr = ctr.getError() + middleErr = ctr.GetError() } } return res, err @@ -1897,7 +2491,7 @@ func TestToolErrorMiddleware(t *testing.T) { t.Fatal("want error, got none") } // Clients can't see the error, because it isn't marshaled. - if err := res.getError(); err != nil { + if err := res.GetError(); err != nil { t.Fatalf("got %v, want nil", err) } if middleErr != errTestFailure { @@ -1905,4 +2499,1454 @@ func TestToolErrorMiddleware(t *testing.T) { } } -var ctrCmpOpts = []cmp.Option{cmp.AllowUnexported(CallToolResult{})} +func TestSetErrorPreservesContent(t *testing.T) { + for _, tt := range []struct { + name string + content []Content + err error + wantContent string + }{ + { + name: "nil content", + err: errors.New("internal failure"), + wantContent: "internal failure", + }, + { + name: "empty slice content", + content: []Content{}, + err: errors.New("internal failure"), + wantContent: "internal failure", + }, + { + name: "existing content preserved", + content: []Content{&TextContent{Text: "user-friendly msg"}}, + err: errors.New("db timeout"), + wantContent: "user-friendly msg", + }, + } { + t.Run(tt.name, func(t *testing.T) { + res := CallToolResult{Content: tt.content} + res.SetError(tt.err) + if !res.IsError { + t.Fatal("want IsError=true") + } + if got := res.Content[0].(*TextContent).Text; got != tt.wantContent { + t.Errorf("Content text = %q, want %q", got, tt.wantContent) + } + if got := res.GetError(); got != tt.err { + t.Errorf("GetError() = %v, want %v", got, tt.err) + } + }) + } +} + +var ctrCmpOpts = []cmp.Option{ + cmpopts.IgnoreUnexported(CallToolResult{}, GetPromptResult{}, ReadResourceResult{}), + // Server responses under the >= 2026-07-28 protocol carry an auto-populated + // [MetaKeyServerInfo] entry; tests that compare result bodies against + // hand-crafted expected values should ignore it. + cmpopts.IgnoreFields(CallToolResult{}, "Meta"), + cmpopts.IgnoreFields(GetPromptResult{}, "Meta"), + cmpopts.IgnoreFields(ReadResourceResult{}, "Meta"), +} + +// runSubscriptionsListenTest exercises the SEP-2575 auto-listen flow end-to-end +// against the supplied transport pair. It captures every notification and the +// acknowledgment the client sees, then asserts: +// +// - the auto-listen issued by Client.Connect is acknowledged with a tagged +// subscription ID; +// - tool and prompt list-changed notifications are delivered to the matching +// handlers, each carrying the same subscription ID as the ack; +// - the subscription persists across multiple unrelated changes; +// - cs.Close() ends the subscription and further changes don't deliver. +func runSubscriptionsListenTest(t *testing.T, client *Client, server *Server, ct Transport, events chan subListenEvent) { + t.Helper() + + ctx, topCancel := context.WithTimeout(context.Background(), 30*time.Second) + defer topCancel() + + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + + waitFor := func(kind string) subListenEvent { + t.Helper() + select { + case e := <-events: + if e.kind != kind { + t.Fatalf("got event %q (id=%s), want kind %q", e.kind, e.id, kind) + } + return e + case <-time.After(5 * time.Second): + t.Fatalf("timed out waiting for %q event", kind) + return subListenEvent{} + } + } + expectNoEvent := func(d time.Duration) { + t.Helper() + select { + case e := <-events: + t.Fatalf("unexpected event %q (id=%s)", e.kind, e.id) + case <-time.After(d): + } + } + + ack := waitFor("ack") + if ack.id == "" { + t.Fatalf("acknowledgment missing subscription ID") + } + + server.AddTool(&Tool{Name: "t2", InputSchema: &jsonschema.Schema{Type: "object"}}, nil) + if e := waitFor("tool"); e.id != ack.id { + t.Errorf("first tool notif id = %s, want %s", e.id, ack.id) + } + + server.AddPrompt(&Prompt{Name: "p2"}, nil) + if e := waitFor("prompt"); e.id != ack.id { + t.Errorf("first prompt notif id = %s, want %s", e.id, ack.id) + } + + server.AddTool(&Tool{Name: "t3", InputSchema: &jsonschema.Schema{Type: "object"}}, nil) + if e := waitFor("tool"); e.id != ack.id { + t.Errorf("second tool notif id = %s, want %s", e.id, ack.id) + } + server.AddPrompt(&Prompt{Name: "p3"}, nil) + if e := waitFor("prompt"); e.id != ack.id { + t.Errorf("second prompt notif id = %s, want %s", e.id, ack.id) + } + expectNoEvent(notificationDelay * 5) + + cs.Close() + time.Sleep(50 * time.Millisecond) + + server.AddTool(&Tool{Name: "t4", InputSchema: &jsonschema.Schema{Type: "object"}}, nil) + server.AddPrompt(&Prompt{Name: "p4"}, nil) + expectNoEvent(notificationDelay * 20) +} + +type subListenEvent struct { + kind string // "ack", "tool", "prompt" + id string // subscription ID from _meta, stringified for cross-encoding equality +} + +// newSubListenClient returns a client wired to push every ack and every +// list-changed notification it receives into events, tagged with the kind +// and the subscription ID extracted from _meta. +func newSubListenClient(events chan subListenEvent) *Client { + asEvent := func(kind string, raw any) subListenEvent { + return subListenEvent{kind, fmt.Sprint(raw)} + } + c := NewClient(testImpl, &ClientOptions{ + ToolListChangedHandler: func(_ context.Context, req *ToolListChangedRequest) { + events <- asEvent("tool", req.Params.Meta[MetaKeySubscriptionID]) + }, + PromptListChangedHandler: func(_ context.Context, req *PromptListChangedRequest) { + events <- asEvent("prompt", req.Params.Meta[MetaKeySubscriptionID]) + }, + }) + c.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == notificationSubscriptionsAck { + if cr, ok := req.(*ClientRequest[*SubscriptionsAcknowledgedParams]); ok && cr.Params != nil { + events <- asEvent("ack", cr.Params.Meta[MetaKeySubscriptionID]) + } + } + return next(ctx, method, req) + } + }) + return c +} + +func newSubListenServer() *Server { + s := NewServer(testImpl, nil) + AddTool(s, &Tool{Name: "t1"}, sayHi) + s.AddPrompt(&Prompt{Name: "p1"}, nil) + return s +} + +// TestSubscriptionsListen_InMemory exercises the listen flow over the +// session-shared in-memory transport (semantically equivalent to STDIO). +// Cancellation here propagates via notifications/cancelled. +func TestSubscriptionsListen_InMemory(t *testing.T) { + events := make(chan subListenEvent, 64) + server := newSubListenServer() + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + runSubscriptionsListenTest(t, newSubListenClient(events), server, ct, events) +} + +// TestSubscriptionsListen_Streamable exercises the listen flow over a +// stateless HTTP server (SEP-2575). Each listen rides its own SSE response +// stream; cs.Close() tears it down. +func TestSubscriptionsListen_Streamable(t *testing.T) { + events := make(chan subListenEvent, 64) + server := newSubListenServer() + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true}, + ) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + runSubscriptionsListenTest(t, newSubListenClient(events), server, + &StreamableClientTransport{Endpoint: httpServer.URL}, events) +} + +// TestSubscriptionsListen_NoHandlersNoListen verifies that a new-protocol +// client without any list-changed handlers registered does not open an +// auto-listen on connect, and therefore does not receive any acknowledgment +// or downstream notifications. +func TestSubscriptionsListen_NoHandlersNoListen(t *testing.T) { + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + events := make(chan subListenEvent, 8) + server := newSubListenServer() + ct, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + + c := NewClient(testImpl, nil) + c.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == notificationSubscriptionsAck { + events <- subListenEvent{"ack", ""} + } + return next(ctx, method, req) + } + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + server.AddTool(&Tool{Name: "t2", InputSchema: &jsonschema.Schema{Type: "object"}}, nil) + + select { + case e := <-events: + t.Fatalf("unexpected event %q on no-handler client", e.kind) + case <-time.After(notificationDelay * 10): + } +} + +// TestSubscriptionsListen_MissingNotifications verifies that a +// subscriptions/listen request without the required "notifications" field +// is rejected with an invalid params error. +func TestSubscriptionsListen_MissingNotifications(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + server := newSubListenServer() + _, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + + // Invoke the server handler directly with a nil Notifications field, + // simulating a request whose params omit the required field on the wire. + // The client-side subscriptionsListen does not await the RPC response + // (the call's lifetime is the notification stream), so we assert the + // server-side behavior at the handler level. + id, err := jsonrpc.MakeID("test-1") + if err != nil { + t.Fatalf("MakeID: %v", err) + } + reqCtx := context.WithValue(ctx, idContextKey{}, id) + req := &SubscriptionsListenRequest{ + Session: ss, + Params: &SubscriptionsListenParams{}, // Notifications is nil + } + _, err = server.subscriptionsListen(reqCtx, req) + if err == nil { + t.Fatal("expected error for missing notifications field, got nil") + } + var jerr *jsonrpc.Error + if !errors.As(err, &jerr) { + t.Fatalf("expected *jsonrpc.Error, got %T: %v", err, err) + } + if jerr.Code != jsonrpc.CodeInvalidParams { + t.Errorf("error code = %d, want %d", jerr.Code, jsonrpc.CodeInvalidParams) + } +} + +// resourceSubServer builds a server that advertises resource subscriptions +// and records every Subscribe/Unsubscribe handler invocation through chans. +func resourceSubServer(t *testing.T, subCh, unsubCh chan string) *Server { + t.Helper() + s := NewServer(testImpl, &ServerOptions{ + SubscribeHandler: func(_ context.Context, r *SubscribeRequest) error { + subCh <- r.Params.URI + return nil + }, + UnsubscribeHandler: func(_ context.Context, r *UnsubscribeRequest) error { + unsubCh <- r.Params.URI + return nil + }, + }) + s.AddResource(&Resource{Name: "r1", URI: "file:///r1"}, nil) + return s +} + +// resourceSubEvent is one delivered notifications/resources/updated. +type resourceSubEvent struct { + uri string + id string // _meta subscription ID, stringified +} + +// TestResourceSubscriptionsSEP2575_Streamable verifies the Subscribe -> +// ResourceUpdated path on a stateless Streamable HTTP server. +// +// Caveat: per-subscription Unsubscribe is intentionally NOT verified here. +// In stateless Streamable HTTP mode the subscriptions/listen handler blocks +// on its request context, and neither the HTTP POST disconnect nor the +// separate notifications/cancelled POST currently propagates to that +// handler's context. The handler only unwinds when the server next attempts +// a write to the (now-dead) SSE stream and the writeErr branch in the +// jsonrpc2 layer cancels the in-flight request. To keep the test +// hermetic we therefore trigger a write at the end by adding a resource, +// which fires notifications/resources/list_changed on the auto-listen path +// (if any) and on the per-URI listen, causing the listen handler to unwind. +// The spec-correct fix is to plumb the POST's request context down to the +// subscriptionsListen handler so HTTP disconnect is observed directly; this +// is tracked separately. +func TestResourceSubscriptions_Streamable(t *testing.T) { + + subCh := make(chan string, 8) + unsubCh := make(chan string, 8) + events := make(chan resourceSubEvent, 16) + + server := resourceSubServer(t, subCh, unsubCh) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true}, + ) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + + c := NewClient(testImpl, &ClientOptions{ + ResourceUpdatedHandler: func(_ context.Context, req *ResourceUpdatedNotificationRequest) { + id := "" + if req.Params != nil && req.Params.Meta != nil { + id = fmt.Sprint(req.Params.Meta[MetaKeySubscriptionID]) + } + events <- resourceSubEvent{uri: req.Params.URI, id: id} + }, + }) + cs, err := c.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, + &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("subscribe r1: %v", err) + } + select { + case got := <-subCh: + if got != "file:///r1" { + t.Fatalf("got URI %q, want %q", got, "file:///r1") + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for SubscribeHandler") + } + + server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r1"}) + select { + case e := <-events: + if e.uri != "file:///r1" { + t.Fatalf("got URI %q, want %q", e.uri, "file:///r1") + } + if e.id == "" || e.id == "" { + t.Fatalf("missing subscription ID on update (got %q)", e.id) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for resource update") + } + + // See test header comment for the explanation of this teardown ritual: + // close the client, then drop server-side TCP, then drive a write that + // will fail (any extra ResourceUpdated for our URI), to unblock the + // in-flight listen handler so httpServer.Close can return. + _ = cs.Close() + httpServer.CloseClientConnections() + server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r1"}) + httpServer.Close() +} + +// TestResourceSubscriptions_InMemory mirrors TestResourceSubscriptions_Streamable +// over an in-memory (stdio-equivalent) transport: the per-URI Subscribe path +// uses notifications/cancelled rather than HTTP disconnect for teardown. +func TestResourceSubscriptions_InMemory(t *testing.T) { + + subCh := make(chan string, 8) + unsubCh := make(chan string, 8) + events := make(chan resourceSubEvent, 16) + + server := resourceSubServer(t, subCh, unsubCh) + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + c := NewClient(testImpl, &ClientOptions{ + ResourceUpdatedHandler: func(_ context.Context, req *ResourceUpdatedNotificationRequest) { + id := "" + if req.Params != nil && req.Params.Meta != nil { + id = fmt.Sprint(req.Params.Meta[MetaKeySubscriptionID]) + } + events <- resourceSubEvent{uri: req.Params.URI, id: id} + }, + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("subscribe r1: %v", err) + } + waitURI := func(ch chan string, want string) { + t.Helper() + select { + case got := <-ch: + if got != want { + t.Fatalf("got URI %q, want %q", got, want) + } + case <-time.After(5 * time.Second): + t.Fatalf("timed out waiting for URI %q", want) + } + } + waitURI(subCh, "file:///r1") + + server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r1"}) + select { + case e := <-events: + if e.uri != "file:///r1" { + t.Fatalf("got URI %q, want %q", e.uri, "file:///r1") + } + if e.id == "" || e.id == "" { + t.Fatalf("missing subscription ID on update (got %q)", e.id) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for resource update") + } + + // On stdio/in-memory, Unsubscribe sends notifications/cancelled which + // reaches the server preempter, cancels the listen handler ctx, and + // triggers the UnsubscribeHandler. + if err := cs.Unsubscribe(ctx, &UnsubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("unsubscribe r1: %v", err) + } + waitURI(unsubCh, "file:///r1") + + // Updates for an unsubscribed URI MUST NOT be delivered. Give the server + // a moment to process the cancellation first. + time.Sleep(50 * time.Millisecond) + server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r1"}) + select { + case e := <-events: + t.Fatalf("unexpected resource update after unsubscribe: %q (id=%s)", e.uri, e.id) + case <-time.After(notificationDelay * 10): + } +} + +// TestResourceSubscriptions_Subscribe_Idempotent verifies that calling +// Subscribe twice for the same URI in the same session is a no-op for the +// second call: it returns nil without invoking SubscribeHandler again and +// without opening a second listen stream. +func TestResourceSubscriptions_Subscribe_Idempotent(t *testing.T) { + + subCh := make(chan string, 8) + unsubCh := make(chan string, 8) + + server := resourceSubServer(t, subCh, unsubCh) + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + c := NewClient(testImpl, &ClientOptions{ + ResourceUpdatedHandler: func(context.Context, *ResourceUpdatedNotificationRequest) {}, + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("first subscribe: %v", err) + } + select { + case got := <-subCh: + if got != "file:///r1" { + t.Fatalf("got URI %q, want %q", got, "file:///r1") + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for first SubscribeHandler") + } + + // Second Subscribe for the same URI returns nil and does NOT fire + // SubscribeHandler again. + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("second subscribe should be no-op, got error: %v", err) + } + select { + case got := <-subCh: + t.Fatalf("duplicate Subscribe should not re-invoke SubscribeHandler (got %q)", got) + case <-time.After(notificationDelay * 10): + } + + // Subsequent Unsubscribe still works (verifies the URI is tracked + // correctly even though the second Subscribe was a no-op). + if err := cs.Unsubscribe(ctx, &UnsubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("unsubscribe: %v", err) + } + select { + case <-unsubCh: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for UnsubscribeHandler") + } +} + +// TestResourceSubscriptions_MultipleURIs verifies that two concurrent +// Subscribe calls on the same session each open their own independent listen +// stream with a distinct subscription ID. Unsubscribing one does not affect +// the other. +func TestResourceSubscriptions_MultipleURIs(t *testing.T) { + + subCh := make(chan string, 8) + unsubCh := make(chan string, 8) + events := make(chan resourceSubEvent, 16) + + server := resourceSubServer(t, subCh, unsubCh) + server.AddResource(&Resource{Name: "r2", URI: "file:///r2"}, nil) + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + c := NewClient(testImpl, &ClientOptions{ + ResourceUpdatedHandler: func(_ context.Context, req *ResourceUpdatedNotificationRequest) { + id := "" + if req.Params != nil && req.Params.Meta != nil { + id = fmt.Sprint(req.Params.Meta[MetaKeySubscriptionID]) + } + events <- resourceSubEvent{uri: req.Params.URI, id: id} + }, + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("subscribe r1: %v", err) + } + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r2"}); err != nil { + t.Fatalf("subscribe r2: %v", err) + } + + // Each Subscribe MUST fire its own SubscribeHandler invocation. + gotURIs := map[string]bool{} + for range 2 { + select { + case got := <-subCh: + gotURIs[got] = true + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for SubscribeHandler") + } + } + if !gotURIs["file:///r1"] || !gotURIs["file:///r2"] { + t.Fatalf("missing SubscribeHandler invocation; got %v", gotURIs) + } + + // Each update delivers exactly one event tagged with that URI's distinct + // subscription ID. + server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r1"}) + ev1 := <-events + server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r2"}) + ev2 := <-events + if ev1.uri != "file:///r1" || ev2.uri != "file:///r2" { + t.Fatalf("got URIs %q and %q, want r1 and r2", ev1.uri, ev2.uri) + } + if ev1.id == "" || ev2.id == "" { + t.Fatalf("missing subscription IDs: r1=%q r2=%q", ev1.id, ev2.id) + } + if ev1.id == ev2.id { + t.Fatalf("r1 and r2 should have distinct subscription IDs, both = %q", ev1.id) + } + + // Unsubscribe r1 only. r2 keeps working. + if err := cs.Unsubscribe(ctx, &UnsubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("unsubscribe r1: %v", err) + } + select { + case got := <-unsubCh: + if got != "file:///r1" { + t.Fatalf("got unsubscribe URI %q, want %q", got, "file:///r1") + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for UnsubscribeHandler") + } + time.Sleep(50 * time.Millisecond) // let server-side cancellation settle + + server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r1"}) + server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r2"}) + + // Only the r2 update should arrive. + select { + case e := <-events: + if e.uri != "file:///r2" || e.id != ev2.id { + t.Fatalf("post-unsubscribe got %q (id=%s), want r2 (id=%s)", e.uri, e.id, ev2.id) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for r2 update after r1 unsubscribe") + } + select { + case e := <-events: + t.Fatalf("unexpected event after r1 unsubscribe: %q (id=%s)", e.uri, e.id) + case <-time.After(notificationDelay * 10): + } + + // Explicit Unsubscribe r2 to verify it still works independently. + if err := cs.Unsubscribe(ctx, &UnsubscribeParams{URI: "file:///r2"}); err != nil { + t.Fatalf("unsubscribe r2: %v", err) + } + select { + case <-unsubCh: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for r2 UnsubscribeHandler") + } +} + +// TestSubscriptionsListen_MultipleSessions verifies that two concurrent +// client sessions on the same server are isolated: a list-changed +// notification is fanned out to BOTH sessions, each with its own subscription +// ID; closing one session does not affect deliveries to the other. +func TestSubscriptionsListen_MultipleSessions(t *testing.T) { + + server := newSubListenServer() + + // Open two clients on the same server. + open := func(t *testing.T) (*ClientSession, chan subListenEvent, *ServerSession) { + t.Helper() + events := make(chan subListenEvent, 16) + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + c := newSubListenClient(events) + cs, err := c.Connect(context.Background(), ct, + &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + return cs, events, ss + } + csA, evA, ssA := open(t) + defer ssA.Close() + csB, evB, ssB := open(t) + defer ssB.Close() + + waitFor := func(ch chan subListenEvent, kind string) subListenEvent { + t.Helper() + select { + case e := <-ch: + if e.kind != kind { + t.Fatalf("got event %q, want %q", e.kind, kind) + } + return e + case <-time.After(5 * time.Second): + t.Fatalf("timed out waiting for %q event", kind) + return subListenEvent{} + } + } + ackA := waitFor(evA, "ack") + ackB := waitFor(evB, "ack") + // Request IDs are per-connection, so both sessions may legitimately use + // the same ID number; we only require that each session's own + // notifications carry its own ack ID. + + // A single change fans out to BOTH sessions, each tagged with that + // session's own ack ID. + server.AddTool(&Tool{Name: "t2", InputSchema: &jsonschema.Schema{Type: "object"}}, nil) + gotA := waitFor(evA, "tool") + gotB := waitFor(evB, "tool") + if gotA.id != ackA.id { + t.Errorf("session A: tool notif id=%s, want %s", gotA.id, ackA.id) + } + if gotB.id != ackB.id { + t.Errorf("session B: tool notif id=%s, want %s", gotB.id, ackB.id) + } + + // Close A; B's subscription must keep delivering. + csA.Close() + time.Sleep(50 * time.Millisecond) + server.AddTool(&Tool{Name: "t3", InputSchema: &jsonschema.Schema{Type: "object"}}, nil) + gotB2 := waitFor(evB, "tool") + if gotB2.id != ackB.id { + t.Errorf("session B after A closed: tool notif id=%s, want %s", gotB2.id, ackB.id) + } + select { + case e := <-evA: + t.Fatalf("session A unexpected event after Close: %q", e.kind) + case <-time.After(notificationDelay * 5): + } + + csB.Close() +} + +// TestSubscriptionsListen_ResourceListChanged covers the resources/list_changed +// auto-listen branch (the other two list-changed types are covered by +// runSubscriptionsListenTest, but resources is not). +func TestSubscriptionsListen_ResourceListChanged(t *testing.T) { + events := make(chan subListenEvent, 16) + + server := NewServer(testImpl, nil) + server.AddResource(&Resource{Name: "r1", URI: "file:///r1"}, nil) + + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + + asEvent := func(kind string, raw any) subListenEvent { + return subListenEvent{kind, fmt.Sprint(raw)} + } + c := NewClient(testImpl, &ClientOptions{ + ResourceListChangedHandler: func(_ context.Context, req *ResourceListChangedRequest) { + id := any(nil) + if req.Params != nil && req.Params.Meta != nil { + id = req.Params.Meta[MetaKeySubscriptionID] + } + events <- asEvent("resource", id) + }, + }) + c.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == notificationSubscriptionsAck { + if cr, ok := req.(*ClientRequest[*SubscriptionsAcknowledgedParams]); ok && cr.Params != nil { + events <- asEvent("ack", cr.Params.Meta[MetaKeySubscriptionID]) + } + } + return next(ctx, method, req) + } + }) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + + waitFor := func(kind string) subListenEvent { + t.Helper() + select { + case e := <-events: + if e.kind != kind { + t.Fatalf("got event %q, want %q", e.kind, kind) + } + return e + case <-time.After(5 * time.Second): + t.Fatalf("timed out waiting for %q", kind) + return subListenEvent{} + } + } + ack := waitFor("ack") + if ack.id == "" { + t.Fatal("acknowledgment missing subscription ID") + } + server.AddResource(&Resource{Name: "r2", URI: "file:///r2"}, nil) + got := waitFor("resource") + if got.id != ack.id { + t.Errorf("resource notif id=%s, want %s", got.id, ack.id) + } + cs.Close() +} + +// TestSubscriptionsListen_DisconnectScrubsMaps verifies that closing a +// session removes its entries from the server's three per-type subscription +// maps via Server.disconnect. +func TestSubscriptionsListen_DisconnectScrubsMaps(t *testing.T) { + events := make(chan subListenEvent, 16) + server := newSubListenServer() + + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + c := newSubListenClient(events) + cs, err := c.Connect(context.Background(), ct, + &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + + // Wait for the auto-listen to actually register on the server side. + select { + case e := <-events: + if e.kind != "ack" { + t.Fatalf("got %q, want ack", e.kind) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for ack") + } + + // Server should now have this session in tools+prompts maps (the fixture + // client opts in to both via newSubListenClient). + server.mu.Lock() + if _, ok := server.toolChangeSubscriptions[ss]; !ok { + server.mu.Unlock() + t.Fatal("session missing from toolChangeSubscriptions before close") + } + if _, ok := server.promptChangeSubscriptions[ss]; !ok { + server.mu.Unlock() + t.Fatal("session missing from promptChangeSubscriptions before close") + } + server.mu.Unlock() + + cs.Close() + ss.Close() + // Closures complete asynchronously on the server side. + deadline := time.Now().Add(5 * time.Second) + for { + server.mu.Lock() + _, inTool := server.toolChangeSubscriptions[ss] + _, inPrompt := server.promptChangeSubscriptions[ss] + _, inResource := server.resourceChangeSubscriptions[ss] + server.mu.Unlock() + if !inTool && !inPrompt && !inResource { + return + } + if time.Now().After(deadline) { + t.Fatalf("subscription maps not scrubbed after Close: tool=%v prompt=%v resource=%v", + inTool, inPrompt, inResource) + } + time.Sleep(10 * time.Millisecond) + } +} + +// newSubListenSubscribeServer is newSubListenServer with a resource and the +// handlers that make the server advertise resources.subscribe, so that a +// listen can carry a resource subscription. +func newSubListenSubscribeServer() *Server { + s := NewServer(testImpl, &ServerOptions{ + SubscribeHandler: func(context.Context, *SubscribeRequest) error { return nil }, + UnsubscribeHandler: func(context.Context, *UnsubscribeRequest) error { return nil }, + }) + AddTool(s, &Tool{Name: "t1"}, sayHi) + s.AddPrompt(&Prompt{Name: "p1"}, nil) + s.AddResource(&Resource{Name: "r1", URI: "file:///r1"}, nil) + return s +} + +func waitSubListenEvent(t *testing.T, events chan subListenEvent, kind string) subListenEvent { + t.Helper() + select { + case e := <-events: + if e.kind != kind { + t.Fatalf("got event %q, want %q", e.kind, kind) + } + return e + case <-time.After(5 * time.Second): + t.Fatalf("timed out waiting for %q", kind) + return subListenEvent{} + } +} + +// TestSubscriptionsListen_TeardownKeepsOtherListens verifies that unwinding +// one listen leaves the other listens on the same session registered. +// Unsubscribe tears down the listen that Subscribe opened for the resource; +// the auto-listen opened by Connect must keep delivering list-changed +// notifications. +func TestSubscriptionsListen_TeardownKeepsOtherListens(t *testing.T) { + events := make(chan subListenEvent, 16) + server := newSubListenSubscribeServer() + + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + cs, err := newSubListenClient(events).Connect(ctx, ct, + &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + autoListen := waitSubListenEvent(t, events, "ack") + + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("Subscribe: %v", err) + } + resourceListen := waitSubListenEvent(t, events, "ack") + if resourceListen.id == autoListen.id { + t.Fatalf("Subscribe reused subscription ID %s", autoListen.id) + } + + if err := cs.Unsubscribe(ctx, &UnsubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("Unsubscribe: %v", err) + } + // Wait for the server to unwind that listen, observed through a registry + // the auto-listen does not appear in. + waitUntil(t, 5*time.Second, "resource listen to unwind", func() bool { + server.mu.Lock() + defer server.mu.Unlock() + _, ok := server.resourceSubscriptions["file:///r1"] + return !ok + }) + + AddTool(server, &Tool{Name: "t2"}, sayHi) + if got := waitSubListenEvent(t, events, "tool"); got.id != autoListen.id { + t.Errorf("tool notification id = %s, want auto-listen %s", got.id, autoListen.id) + } +} + +// TestSubscriptionsListen_TeardownRetiresOwnRegistration verifies that a +// listen still retires what it registered, so that per-listen tracking does +// not turn into a leak. +func TestSubscriptionsListen_TeardownRetiresOwnRegistration(t *testing.T) { + events := make(chan subListenEvent, 8) + server := newSubListenServer() + + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + + // A client with no list-changed handlers does not auto-listen, so the + // listen opened below is the session's only one. + c := NewClient(testImpl, nil) + c.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == notificationSubscriptionsAck { + events <- subListenEvent{"ack", ""} + } + return next(ctx, method, req) + } + }) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + listenCtx, cancelListen := context.WithCancel(context.Background()) + defer cancelListen() + if err := cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{ + Notifications: &NotificationSubscriptions{ToolsListChanged: true, PromptsListChanged: true}, + }); err != nil { + t.Fatalf("listen: %v", err) + } + waitSubListenEvent(t, events, "ack") + + server.mu.Lock() + _, inTool := server.toolChangeSubscriptions[ss] + _, inPrompt := server.promptChangeSubscriptions[ss] + server.mu.Unlock() + if !inTool || !inPrompt { + t.Fatal("listen not registered") + } + + cancelListen() + waitUntil(t, 5*time.Second, "registrations to be retired", func() bool { + server.mu.Lock() + defer server.mu.Unlock() + _, inTool := server.toolChangeSubscriptions[ss] + _, inPrompt := server.promptChangeSubscriptions[ss] + return !inTool && !inPrompt + }) +} + +// waitUntil polls cond until it reports true, failing the test after timeout. +func waitUntil(t *testing.T, timeout time.Duration, what string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(timeout) + for !cond() { + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %s", what) + } + time.Sleep(10 * time.Millisecond) + } +} + +// TestSubscriptionsListen_RespectsServerCapabilities verifies that during +// Connect the client only opens a SEP-2575 subscriptions/listen stream for the +// change notifications the server advertised during capability negotiation. +func TestSubscriptionsListen_RespectsServerCapabilities(t *testing.T) { + for _, tc := range []struct { + name string + caps *ServerCapabilities + wantListen bool + }{ + { + name: "advertised", + caps: &ServerCapabilities{Tools: &ToolCapabilities{ListChanged: true}}, + wantListen: true, + }, + { + name: "not advertised", + caps: &ServerCapabilities{Tools: &ToolCapabilities{}}, + wantListen: false, + }, + } { + t.Run(tc.name, func(t *testing.T) { + server := NewServer(testImpl, &ServerOptions{Capabilities: tc.caps}) + AddTool(server, &Tool{Name: "t1"}, sayHi) + + listenSeen := make(chan struct{}, 1) + server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == methodSubscriptionsListen { + select { + case listenSeen <- struct{}{}: + default: + } + } + return next(ctx, method, req) + } + }) + + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + + c := NewClient(testImpl, &ClientOptions{ + ToolListChangedHandler: func(context.Context, *ToolListChangedRequest) {}, + }) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + if _, err := cs.ListTools(ctx, nil); err != nil { + t.Fatalf("ListTools: %v", err) + } + + if tc.wantListen { + select { + case <-listenSeen: + case <-time.After(5 * time.Second): + t.Fatal("expected subscriptions/listen, but none was sent") + } + return + } + + select { + case <-listenSeen: + t.Fatal("client sent subscriptions/listen for an unadvertised capability") + case <-time.After(notificationDelay * 20): + } + }) + } +} + +// TestServerSessionCloseWithActiveListen is a regression test for +// modelcontextprotocol/go-sdk#1160: ServerSession.Close must not deadlock +// when the client has an active subscriptions/listen stream. Previously, +// the server-side handler parked on ctx.Done and Close waited forever for +// the in-flight request to drain. +// +// The client's auto-listen (triggered by registering a list-changed handler) +// opens the stream on Connect — no explicit Subscribe is needed. The server +// must expose the corresponding list-changed capability, which happens +// automatically when at least one tool/prompt/resource is registered. +func TestServerSessionCloseWithActiveListen(t *testing.T) { + ctx := context.Background() + s := NewServer(&Implementation{Name: "s", Version: "0"}, nil) + AddTool(s, &Tool{Name: "t"}, sayHi) + + ct, st := NewInMemoryTransports() + if _, err := s.Connect(ctx, st, nil); err != nil { + t.Fatalf("server connect: %v", err) + } + c := NewClient(&Implementation{Name: "c", Version: "0"}, &ClientOptions{ + ToolListChangedHandler: func(context.Context, *ToolListChangedRequest) {}, + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + // Give the auto-listen request time to reach the server-side handler. + time.Sleep(20 * time.Millisecond) + + var ss *ServerSession + for x := range s.Sessions() { + ss = x + break + } + if ss == nil { + t.Fatal("no server session found") + } + + // Sanity check: the auto-listen must actually have registered an entry in + // listenIDs, otherwise the test below would trivially pass without + // exercising the fix. + if n := listenIDsCount(ss); n == 0 { + t.Fatal("expected auto-listen to register a request ID on the server session") + } + + done := make(chan error, 1) + go func() { done <- ss.Close() }() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("ServerSession.Close deadlocked with an active subscriptions/listen") + } +} + +// listenIDsCount reports how many request IDs are currently recorded in the +// session's listenIDs set. +func listenIDsCount(ss *ServerSession) int { + ss.mu.Lock() + defer ss.mu.Unlock() + return len(ss.listenIDs) +} + +// A completed listen stream must not leave a stale entry in the session. +func TestListenPrunedAfterSingleCompletion(t *testing.T) { + cs, ss, cleanup := basicClientServerConnection(t, nil, nil, func(s *Server) { + AddTool(s, &Tool{Name: "t"}, sayHi) + }) + _ = cleanup + + ctx := context.Background() + lctx, cancel := context.WithCancel(ctx) + go cs.subscriptionsListen(lctx, &SubscriptionsListenParams{ + Notifications: &NotificationSubscriptions{ToolsListChanged: true}, + }) + time.Sleep(30 * time.Millisecond) + cancel() // peer cancels: server handler returns + time.Sleep(30 * time.Millisecond) + + if n := listenIDsCount(ss); n != 0 { + t.Fatalf("completed listen left %d stale entry/ies", n) + } +} + +// Completed listens must not accumulate: the slice grows without bound today. +func TestListenIDsDoNotAccumulate(t *testing.T) { + cs, ss, cleanup := basicClientServerConnection(t, nil, nil, func(s *Server) { + AddTool(s, &Tool{Name: "t"}, sayHi) + }) + _ = cleanup + + ctx := context.Background() + + const cycles = 15 + for range cycles { + lctx, cancel := context.WithCancel(ctx) + go cs.subscriptionsListen(lctx, &SubscriptionsListenParams{ + Notifications: &NotificationSubscriptions{ToolsListChanged: true}, + }) + time.Sleep(15 * time.Millisecond) + cancel() + time.Sleep(10 * time.Millisecond) + } + time.Sleep(30 * time.Millisecond) + + if n := listenIDsCount(ss); n != 0 { + t.Fatalf("listenIDs grew unbounded: %d stale entries after %d completed listens", n, cycles) + } +} + +// The leak is reachable through the public Subscribe/Unsubscribe API: every +// subscription opens a listen stream and unsubscribing completes it, so a real +// client cycling subscriptions on a long-lived session leaks one entry per cycle. +func TestListenIDsLeakViaPublicSubscribeUnsubscribe(t *testing.T) { + cs, ss, cleanup := basicClientServerConnection(t, nil, nil, func(s *Server) { + AddTool(s, &Tool{Name: "t"}, sayHi) + }) + _ = cleanup + + ctx := context.Background() + + const cycles = 3 + for i := range cycles { + uri := fmt.Sprintf("resource://cycle-%d", i) + if err := cs.Subscribe(ctx, &SubscribeParams{URI: uri}); err != nil { + t.Fatalf("Subscribe %d: %v", i, err) + } + time.Sleep(15 * time.Millisecond) + if err := cs.Unsubscribe(ctx, &UnsubscribeParams{URI: uri}); err != nil { + t.Fatalf("Unsubscribe %d: %v", i, err) + } + } + time.Sleep(30 * time.Millisecond) + + if n := listenIDsCount(ss); n != 0 { + t.Fatalf("public Subscribe/Unsubscribe leaked %d stale listenIDs after %d cycles", n, cycles) + } +} + +func TestCustomMethods(t *testing.T) { + type searchParams struct { + ParamsBase + Query string `json:"query"` + Limit int `json:"limit,omitempty"` + } + + type searchResult struct { + ResultBase + Hits []string `json:"hits"` + Total int `json:"total"` + } + + callCustom := func(ctx context.Context, conn *jsonrpc2.Connection, method string, params, result any) error { + return conn.Call(ctx, method, params).Await(ctx, result) + } + + ctx := context.Background() + s := NewServer(testImpl, nil) + + if err := AddReceivingCustomMethod(s, "acme/search", func(ctx context.Context, ss *ServerSession, params *searchParams) (*searchResult, error) { + hits := []string{"result for " + params.Query} + return &searchResult{ + Hits: hits, + Total: len(hits), + }, nil + }); err != nil { + t.Fatal(err) + } + + ct, st := NewInMemoryTransports() + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ss.Close() }) + + c := NewClient(testImpl, nil) + if err := AddSendingCustomMethod[*searchParams, *searchResult](c, "acme/search"); err != nil { + t.Fatal(err) + } + + cs, err := c.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = cs.Close() }) + + // Test raw JSON-RPC call. + var result1 searchResult + if err := callCustom(ctx, cs.getConn(), "acme/search", &searchParams{Query: "hello", Limit: 10}, &result1); err != nil { + t.Fatal(err) + } + if len(result1.Hits) != 1 || result1.Hits[0] != "result for hello" { + t.Errorf("raw call: unexpected hits: %v", result1.Hits) + } + if result1.Total != 1 { + t.Errorf("raw call: unexpected total: %d", result1.Total) + } + + // Test the typed CallCustomMethod helper. + result2, err := CallCustomMethod[*searchParams, *searchResult]( + ctx, cs, "acme/search", &searchParams{Query: "world"}) + if err != nil { + t.Fatal(err) + } + if len(result2.Hits) != 1 || result2.Hits[0] != "result for world" { + t.Errorf("CallCustomMethod: unexpected hits: %v", result2.Hits) + } + if result2.Total != 1 { + t.Errorf("CallCustomMethod: unexpected total: %d", result2.Total) + } + + // CallCustomMethod must reject methods that were never registered. + if _, err := CallCustomMethod[*searchParams, *searchResult]( + ctx, cs, "acme/unregistered", &searchParams{Query: "x"}); err == nil { + t.Error("CallCustomMethod: expected error for unregistered method, got nil") + } +} + +func TestAddCustomMethodRejectsStandardMethods(t *testing.T) { + type params struct{ ParamsBase } + type result struct{ ResultBase } + + t.Run("server", func(t *testing.T) { + s := NewServer(testImpl, nil) + err := AddReceivingCustomMethod(s, "tools/call", + func(ctx context.Context, ss *ServerSession, p *params) (*result, error) { + return &result{}, nil + }) + if err == nil { + t.Fatal("AddReceivingCustomMethod: expected error when shadowing a standard method, got nil") + } + }) + + t.Run("client", func(t *testing.T) { + c := NewClient(testImpl, nil) + if err := AddSendingCustomMethod[*params, *result](c, "tools/call"); err == nil { + t.Fatal("AddSendingCustomMethod: expected error when shadowing a standard method, got nil") + } + }) +} + +// TestCallCustomMethodTypedNilParams exercises the typed-nil params path. +// User param types embed ParamsBase, so the inherited isNil forwarder would +// dereference a typed-nil outer if injectRequestMeta were called with it. +// CallCustomMethod must allocate a fresh value before the meta-injection step. +func TestCallCustomMethodTypedNilParams(t *testing.T) { + type pingParams struct{ ParamsBase } + type pingResult struct{ ResultBase } + + ctx := context.Background() + s := NewServer(testImpl, nil) + if err := AddReceivingCustomMethod(s, "acme/ping", + func(ctx context.Context, ss *ServerSession, p *pingParams) (*pingResult, error) { + return &pingResult{}, nil + }); err != nil { + t.Fatal(err) + } + ct, st := NewInMemoryTransports() + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ss.Close() }) + + c := NewClient(testImpl, nil) + if err := AddSendingCustomMethod[*pingParams, *pingResult](c, "acme/ping"); err != nil { + t.Fatal(err) + } + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = cs.Close() }) + + var typedNil *pingParams + if _, err := CallCustomMethod[*pingParams, *pingResult](ctx, cs, "acme/ping", typedNil); err != nil { + t.Fatalf("CallCustomMethod with typed-nil params: %v", err) + } +} + +func TestServerLogLevelDoesNotLeakBetweenNewProtocolRequests(t *testing.T) { + ctx := context.Background() + s := NewServer(testImpl, nil) + _, st := NewInMemoryTransports() + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ss.Close() }) + + logged := make(chan LoggingLevel, 1) + s.AddSendingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == notificationLoggingMessage { + logged <- req.GetParams().(*LoggingMessageParams).Level + return nil, nil + } + return next(ctx, method, req) + } + }) + + started := make(chan struct{}) + release := make(chan struct{}) + AddTool(s, &Tool{Name: "blocked-log"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { + close(started) + <-release + if err := req.Session.Log(ctx, &LoggingMessageParams{Level: "warning", Data: "request log"}); err != nil { + return nil, nil, err + } + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + AddTool(s, &Tool{Name: "noop"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + + withLogLevel := &CallToolParams{Name: "blocked-log"} + withLogLevel.SetMeta(newProtocolMeta("warning")) + errc := make(chan error, 1) + go func() { + _, err := ss.handle(ctx, req(1, methodCallTool, withLogLevel)) + errc <- err + }() + + <-started + withoutLogLevel := &CallToolParams{Name: "noop"} + withoutLogLevel.SetMeta(newProtocolMeta("")) + if _, err := ss.handle(ctx, req(2, methodCallTool, withoutLogLevel)); err != nil { + t.Fatal(err) + } + close(release) + if err := <-errc; err != nil { + t.Fatal(err) + } + + select { + case got := <-logged: + if got != "warning" { + t.Fatalf("logged level = %q, want warning", got) + } + default: + t.Fatal("request-scoped warning log was suppressed after another request cleared session log level") + } +} + +func newProtocolMeta(logLevel LoggingLevel) Meta { + m := Meta{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: testImpl, + MetaKeyClientCapabilities: (&ClientCapabilities{}).toV2(), + } + if logLevel != "" { + m[MetaKeyLogLevel] = logLevel + } + return m +} diff --git a/mcp/mrtr.go b/mcp/mrtr.go new file mode 100644 index 0000000..fdf98ec --- /dev/null +++ b/mcp/mrtr.go @@ -0,0 +1,299 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package mcp + +import ( + "context" + "fmt" + "log/slog" + "sync" + + "github.com/modelcontextprotocol/go-sdk/jsonrpc" + "golang.org/x/sync/errgroup" +) + +const maxMultiRoundTripRetries = 10 +const maxLoadSheddingMultiRoundTripRetries = 3 + +// MultiRoundTripOptions configures the client-side multi round-trip request (SEP-2322) +// middleware. The middleware is enabled by default and automatically fulfills input +// requests from the server by invoking the appropriate client handlers and +// retrying the original call. +type MultiRoundTripOptions struct { + // Disabled prevents the automatic multi-round-tirp middleware from being installed. + // When true, the client returns input-required results directly and callers must + // handle the retry loop themselves using [CallToolResult.NeedsInput], + // [GetPromptResult.NeedsInput], or [ReadResourceResult.NeedsInput]. + Disabled bool +} + +type multiRoundTripResponse interface { + setResultType(resultType) + inputRequests() map[string]InputRequest + requestState() string + hasContent() bool +} + +// validateMultiRoundTripResult rejects a result that carries both content and +// input requests. [annotateResultType] labels the result later, in +// [ServerSession.handle], so this function does not set resultType. +func validateMultiRoundTripResult(logger *slog.Logger, res multiRoundTripResponse) error { + if res == nil { + return nil + } + if res.inputRequests() != nil && res.hasContent() { + logger.Warn("handler returned both content and inputRequests") + return &jsonrpc.Error{ + Code: jsonrpc.CodeInternalError, + Message: "server bug: result has both content and inputRequests", + } + } + return nil +} + +func clientSupportsMultiRoundTrip(ss *ServerSession) bool { + protocolVersion := latestProtocolVersion + if iparams := ss.InitializeParams(); iparams != nil { + protocolVersion = iparams.ProtocolVersion + } + return protocolVersion >= protocolVersion20260728 +} + +func clientMultiRoundTripMiddleware() Middleware { + return func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method != methodCallTool && method != methodGetPrompt && method != methodReadResource { + return next(ctx, method, req) + } + + loadSheddingFailures := 0 + for retries := 1; ; retries++ { + res, err := next(ctx, method, req) + if err != nil { + return res, err + } + mrtrResult, ok := res.(multiRoundTripResponse) + if !ok { + return res, nil + } + reqMap := mrtrResult.inputRequests() + if reqMap == nil { + return res, nil + } + if len(reqMap) == 0 { + loadSheddingFailures++ + } + if loadSheddingFailures >= maxLoadSheddingMultiRoundTripRetries { + return nil, fmt.Errorf("multi-round-trip: exceeded maximum load-shedding retries (%d)", maxLoadSheddingMultiRoundTripRetries) + } + if retries >= maxMultiRoundTripRetries { + return nil, fmt.Errorf("multi-round-trip: exceeded maximum retries (%d)", maxMultiRoundTripRetries) + } + cs, ok := req.GetSession().(*ClientSession) + if !ok { + return res, nil + } + responses, err := fulfillInputRequests(ctx, cs, reqMap) + if err != nil { + return nil, err + } + req, err = setMultiRoundTripRetryParams(method, req, responses, mrtrResult.requestState()) + if err != nil { + return nil, err + } + } + } + } +} + +// serverMultiRoundTripMiddleware is a receiving middleware for servers that transparently +// handles multi-round-trip for clients on older protocol versions. When a handler returns +// InputRequests and the client does not support multi-round-trip, the middleware fulfills +// the requests by calling the client directly and reinvokes the handler once with the responses. +func serverMultiRoundTripMiddleware() Middleware { + return func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method != methodCallTool && method != methodGetPrompt && method != methodReadResource { + return next(ctx, method, req) + } + + ss, ok := req.GetSession().(*ServerSession) + if !ok { + return next(ctx, method, req) + } + if clientSupportsMultiRoundTrip(ss) { + return next(ctx, method, req) + } + + res, err := next(ctx, method, req) + if err != nil { + return res, err + } + mrtrResult, ok := res.(multiRoundTripResponse) + if !ok { + return res, nil + } + reqMap := mrtrResult.inputRequests() + if reqMap == nil { + return res, nil + } + if len(reqMap) == 0 { + return nil, fmt.Errorf("the server is busy, retry later") + } + responses, err := fulfillServerInputRequests(ctx, ss, reqMap) + if err != nil { + return nil, err + } + req, err = setMultiRoundTripRetryParams(method, req, responses, mrtrResult.requestState()) + if err != nil { + return nil, err + } + return next(ctx, method, req) + } + } +} + +func fulfillServerInputRequests(ctx context.Context, ss *ServerSession, requests InputRequestMap) (InputResponseMap, error) { + g, ctx := errgroup.WithContext(ctx) + var mu sync.Mutex + responses := make(InputResponseMap, len(requests)) + for id, ir := range requests { + g.Go(func() error { + resp, err := fulfillServerInputRequest(ctx, ss, ir) + if err != nil { + return fmt.Errorf("fulfilling input request %q: %w", id, err) + } + mu.Lock() + responses[id] = resp + mu.Unlock() + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, fmt.Errorf("multi-round-trip: %w", err) + } + return responses, nil +} + +func fulfillServerInputRequest(ctx context.Context, ss *ServerSession, ir InputRequest) (InputResponse, error) { + switch p := ir.(type) { + case *ElicitParams: + return ss.Elicit(ctx, p) + case *CreateMessageParams: + return ss.CreateMessageWithTools(ctx, createMessageParamsToWithTools(p)) + case *CreateMessageWithToolsParams: + return ss.CreateMessageWithTools(ctx, p) + case *ListRootsParams: + return ss.ListRoots(ctx, p) + default: + return nil, fmt.Errorf("unknown input request type: %T", ir) + } +} + +func createMessageParamsToWithTools(p *CreateMessageParams) *CreateMessageWithToolsParams { + var msgs []*SamplingMessageV2 + for _, m := range p.Messages { + msgs = append(msgs, &SamplingMessageV2{Content: []Content{m.Content}, Role: m.Role}) + } + return &CreateMessageWithToolsParams{ + Meta: p.Meta, + IncludeContext: p.IncludeContext, + MaxTokens: p.MaxTokens, + Messages: msgs, + Metadata: p.Metadata, + ModelPreferences: p.ModelPreferences, + StopSequences: p.StopSequences, + SystemPrompt: p.SystemPrompt, + Temperature: p.Temperature, + } +} + +// setMultiRoundTripRetryParams returns the request for the next round of a +// multi-round-trip retry: a request whose params are a shallow copy of the +// original, carrying the fulfilled responses and the request state to echo. +// The original params are never mutated: on the client side they are +// caller-owned, and a params struct reused across calls must not carry one +// call's inputResponses and requestState into the next, where a server +// gating on them (for example an elicitation confirmation) would treat the +// new call as already answered. +// +// An unrecognized params type is an error. It cannot occur on the server +// side, but client sending middleware runs before the retry loop and may +// substitute its own params. +func setMultiRoundTripRetryParams(method string, req Request, responses InputResponseMap, state string) (Request, error) { + var retryParams Params + switch p := req.GetParams().(type) { + case *CallToolParams: + cp := *p + cp.InputResponses = responses + cp.RequestState = state + retryParams = &cp + case *CallToolParamsRaw: + cp := *p + cp.InputResponses = responses + cp.RequestState = state + retryParams = &cp + case *GetPromptParams: + cp := *p + cp.InputResponses = responses + cp.RequestState = state + retryParams = &cp + case *ReadResourceParams: + cp := *p + cp.InputResponses = responses + cp.RequestState = state + retryParams = &cp + default: + return nil, fmt.Errorf("multi-round-trip: unsupported params type %T", p) + } + switch s := req.GetSession().(type) { + case *ClientSession: + return newClientRequest(s, retryParams), nil + case *ServerSession: + // Rebuild with the same constructor handleReceive uses: the + // receiving dispatch type-asserts on the request's concrete + // *ServerRequest[P] type, and Extra must carry over. + return s.receivingMethodInfos()[method].newRequest(s, retryParams, req.GetExtra()), nil + default: + return nil, fmt.Errorf("multi-round-trip: unsupported session type %T", s) + } +} + +func fulfillInputRequests(ctx context.Context, cs *ClientSession, requests InputRequestMap) (InputResponseMap, error) { + g, ctx := errgroup.WithContext(ctx) + var mu sync.Mutex + responses := make(InputResponseMap, len(requests)) + for id, ir := range requests { + g.Go(func() error { + resp, err := fulfillInputRequest(ctx, cs, ir) + if err != nil { + return fmt.Errorf("fulfilling input request %q: %w", id, err) + } + mu.Lock() + responses[id] = resp + mu.Unlock() + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, fmt.Errorf("multi round-trip: %w", err) + } + return responses, nil +} + +func fulfillInputRequest(ctx context.Context, cs *ClientSession, ir InputRequest) (InputResponse, error) { + switch p := ir.(type) { + case *ElicitParams: + return cs.client.elicit(ctx, newClientRequest(cs, p)) + case *CreateMessageParams: + return cs.client.createMessage(ctx, &CreateMessageWithToolsRequest{Session: cs, Params: createMessageParamsToWithTools(p)}) + case *CreateMessageWithToolsParams: + return cs.client.createMessage(ctx, &CreateMessageWithToolsRequest{Session: cs, Params: p}) + case *ListRootsParams: + return cs.client.listRoots(ctx, newClientRequest(cs, p)) + default: + return nil, fmt.Errorf("unknown input request type: %T", ir) + } +} diff --git a/mcp/mrtr_test.go b/mcp/mrtr_test.go new file mode 100644 index 0000000..d966219 --- /dev/null +++ b/mcp/mrtr_test.go @@ -0,0 +1,913 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +//lint:file-ignore SA1019 tests exercise deprecated SEP-2577 APIs (roots, sampling, logging). + +package mcp + +import ( + "context" + "fmt" + "slices" + "strings" + "sync/atomic" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/google/jsonschema-go/jsonschema" +) + +func TestMultiRoundTrip_ManualRetry(t *testing.T) { + type deployResult struct { + Deployed bool `json:"deployed"` + Reason string `json:"reason,omitempty"` + } + + ctx := context.Background() + + srv := NewServer(testImpl, nil) + AddTool(srv, &Tool{Name: "deploy"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, *deployResult, error) { + if len(req.Params.InputResponses) == 0 { + return &CallToolResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Deploy to production?"}}, + RequestState: "deployment-123", + }, nil, nil + } + + resp, ok := req.Params.InputResponses["confirm"] + if !ok { + return &CallToolResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Please confirm (retry)"}}, + }, nil, nil + } + + if req.Params.RequestState == "" { + return &CallToolResult{}, &deployResult{Deployed: false, Reason: "no_state"}, nil + } + if elicitResult := resp.(*ElicitResult); elicitResult != nil && elicitResult.Action != "accept" { + return &CallToolResult{}, &deployResult{Deployed: false, Reason: "cancelled"}, nil + } + + return &CallToolResult{}, &deployResult{Deployed: true}, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + MultiRoundTrip: &MultiRoundTripOptions{Disabled: true}, + }) + + // Round 1: initiate deployment + res, err := conn.CallTool(ctx, &CallToolParams{Name: "deploy"}) + if err != nil { + t.Fatalf("CallTool() error = %v", err) + } + if !res.NeedsInput() { + t.Fatal("NeedsInput() = false, want true") + } + if got := len(res.InputRequests); got != 1 { + t.Fatalf("len(res.InputRequests) = %d, want 1", got) + } + if _, ok := res.InputRequests["confirm"].(*ElicitParams); !ok { + t.Fatalf("res.InputRequests[confirm] type = %T, want *ElicitParams", res.InputRequests["confirm"]) + } + + // Round 2: retry with confirmation + res, err = conn.CallTool(ctx, &CallToolParams{ + Name: "deploy", + InputResponses: InputResponseMap{ + "confirm": &ElicitResult{Action: "accept", Content: map[string]any{"ok": true}}, + }, + RequestState: res.RequestState, + }) + if err != nil { + t.Fatalf("CallTool() follow-up error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after follow-up, want false") + } + + if diff := cmp.Diff(map[string]any{"deployed": true}, res.StructuredContent, ctrCmpOpts...); diff != "" { + t.Errorf("result mismatch (-want +got):\n%s", diff) + } +} + +func TestMultiRoundTrip_AutoRetry(t *testing.T) { + + tests := []struct { + name string + inputRequests InputRequestMap + wantResult map[string]any + }{ + { + name: "elicit", + inputRequests: InputRequestMap{ + "confirm": &ElicitParams{Message: "Deploy?"}, + }, + wantResult: map[string]any{"ids": []any{"confirm"}}, + }, + { + name: "createMessage", + inputRequests: InputRequestMap{ + "summarize": &CreateMessageParams{ + Messages: []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "summarize"}}}, + MaxTokens: 100, + }, + }, + wantResult: map[string]any{"ids": []any{"summarize"}}, + }, + { + name: "listRoots", + inputRequests: InputRequestMap{ + "roots": &ListRootsParams{}, + }, + wantResult: map[string]any{"ids": []any{"roots"}}, + }, + { + name: "all three", + inputRequests: InputRequestMap{ + "confirm": &ElicitParams{Message: "OK?"}, + "draft": &CreateMessageParams{ + Messages: []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "write"}}}, + MaxTokens: 50, + }, + "roots": &ListRootsParams{}, + }, + wantResult: map[string]any{"ids": []any{"confirm", "draft", "roots"}}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + inputRequests := tt.inputRequests + AddTool(srv, &Tool{Name: "act"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &CallToolResult{ + InputRequests: inputRequests, + RequestState: "state-1", + }, nil, nil + } + // Collect the IDs of fulfilled responses. + var ids []string + for id := range req.Params.InputResponses { + ids = append(ids, id) + } + slices.Sort(ids) + return &CallToolResult{}, map[string]any{"ids": ids}, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + ElicitationHandler: func(_ context.Context, req *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + CreateMessageHandler: func(_ context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) { + return &CreateMessageResult{ + Model: "test-model", + Role: "assistant", + Content: &TextContent{Text: "response"}, + }, nil + }, + }) + conn.client.AddRoots(&Root{URI: "file:///workspace", Name: "workspace"}) + + res, err := conn.CallTool(ctx, &CallToolParams{Name: "act"}) + if err != nil { + t.Fatalf("CallTool() error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after auto-retry, want false") + } + + // Sort the expected IDs for stable comparison. + if wantIDs, ok := tt.wantResult["ids"].([]any); ok { + slices.SortFunc(wantIDs, func(a, b any) int { + if a.(string) < b.(string) { + return -1 + } + return 1 + }) + } + + if diff := cmp.Diff(tt.wantResult, res.StructuredContent, ctrCmpOpts...); diff != "" { + t.Errorf("result mismatch (-want +got):\n%s", diff) + } + }) + } +} + +// TestMultiRoundTrip_AutoRetryDoesNotMutateCallerParams verifies that +// the client middleware carries inputResponses and requestState on a +// copy of the caller's params: after CallTool returns, the caller's +// struct is unchanged, and reusing it for a second call fulfills the +// input requests again instead of silently replaying the first call's +// answers against a server that gates on them. +func TestMultiRoundTrip_AutoRetryDoesNotMutateCallerParams(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + AddTool(srv, &Tool{Name: "act"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + return &CallToolResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Sure?"}}, + RequestState: "state-1", + }, nil, nil + } + return &CallToolResult{}, map[string]any{"ok": true}, nil + }) + + var elicitations atomic.Int32 + conn := mustConnect(t, srv, &ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + elicitations.Add(1) + return &ElicitResult{Action: "accept"}, nil + }, + }) + + params := &CallToolParams{Name: "act"} + res, err := conn.CallTool(ctx, params) + if err != nil { + t.Fatalf("CallTool() error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after auto-retry, want false") + } + if params.InputResponses != nil { + t.Errorf("params.InputResponses = %v after CallTool, want nil (caller params must not be mutated)", params.InputResponses) + } + if params.RequestState != "" { + t.Errorf("params.RequestState = %q after CallTool, want empty (caller params must not be mutated)", params.RequestState) + } + if got := elicitations.Load(); got != 1 { + t.Fatalf("elicitations = %d, want 1", got) + } + + // Reusing the same params struct must fulfill the input requests + // again, not replay the first call's responses. + if _, err := conn.CallTool(ctx, params); err != nil { + t.Fatalf("CallTool() reuse error = %v", err) + } + if got := elicitations.Load(); got != 2 { + t.Errorf("elicitations after reuse = %d, want 2 (stale responses must not be replayed)", got) + } +} + +// TestMultiRoundTrip_GetPrompt_AutoRetryDoesNotMutateCallerParams is the +// GetPrompt analogue of the CallTool caller-params test. +func TestMultiRoundTrip_GetPrompt_AutoRetryDoesNotMutateCallerParams(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddPrompt(&Prompt{Name: "review"}, func(_ context.Context, req *GetPromptRequest) (*GetPromptResult, error) { + if len(req.Params.InputResponses) == 0 { + return &GetPromptResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Include sensitive data?"}}, + RequestState: "prompt-state", + }, nil + } + return &GetPromptResult{ + Messages: []*PromptMessage{{Role: "user", Content: &TextContent{Text: "review this code"}}}, + }, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + + params := &GetPromptParams{Name: "review"} + res, err := conn.GetPrompt(ctx, params) + if err != nil { + t.Fatalf("GetPrompt() error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after auto-retry, want false") + } + if params.InputResponses != nil { + t.Errorf("params.InputResponses = %v after GetPrompt, want nil (caller params must not be mutated)", params.InputResponses) + } + if params.RequestState != "" { + t.Errorf("params.RequestState = %q after GetPrompt, want empty (caller params must not be mutated)", params.RequestState) + } +} + +// TestMultiRoundTrip_ReadResource_AutoRetryDoesNotMutateCallerParams is +// the ReadResource analogue of the CallTool caller-params test. +func TestMultiRoundTrip_ReadResource_AutoRetryDoesNotMutateCallerParams(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddResource(&Resource{URI: "test://data", Name: "data"}, func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + if len(req.Params.InputResponses) == 0 { + return &ReadResourceResult{ + InputRequests: InputRequestMap{"auth": &ElicitParams{Message: "Authenticate?"}}, + RequestState: "resource-state", + }, nil + } + return &ReadResourceResult{ + Contents: []*ResourceContents{{URI: "test://data", Text: "resource data"}}, + }, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + + params := &ReadResourceParams{URI: "test://data"} + res, err := conn.ReadResource(ctx, params) + if err != nil { + t.Fatalf("ReadResource() error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after auto-retry, want false") + } + if params.InputResponses != nil { + t.Errorf("params.InputResponses = %v after ReadResource, want nil (caller params must not be mutated)", params.InputResponses) + } + if params.RequestState != "" { + t.Errorf("params.RequestState = %q after ReadResource, want empty (caller params must not be mutated)", params.RequestState) + } +} + +// wrappedCallToolParams embeds *CallToolParams so it satisfies Params and +// marshals identically, while being a concrete type the multi-round-trip +// retry does not recognize. +type wrappedCallToolParams struct{ *CallToolParams } + +// TestMultiRoundTrip_UnsupportedRetryParamsType verifies that the client +// retry loop reports an explicit error when the params carry a type it +// cannot build retry params for, instead of silently resending the request +// unchanged until the retry cap. +func TestMultiRoundTrip_UnsupportedRetryParamsType(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + AddTool(srv, &Tool{Name: "act"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, any, error) { + return &CallToolResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Sure?"}}, + RequestState: "state-1", + }, nil, nil + }) + + st, ct := NewInMemoryTransports() + ss, err := srv.Connect(t.Context(), st, nil) + if err != nil { + t.Fatalf("server.Connect() error = %v", err) + } + t.Cleanup(func() { _ = ss.Close() }) + + c := NewClient(testImpl, &ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + // Sending middleware added here wraps outside the multi-round-trip + // middleware installed at client construction, so the retry loop sees + // the substituted params type. + c.AddSendingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == methodCallTool { + cs := req.GetSession().(*ClientSession) + req = newClientRequest[Params](cs, wrappedCallToolParams{req.GetParams().(*CallToolParams)}) + } + return next(ctx, method, req) + } + }) + cs, err := c.Connect(t.Context(), ct, nil) + if err != nil { + t.Fatalf("client.Connect() error = %v", err) + } + t.Cleanup(func() { _ = cs.Close() }) + + _, err = cs.CallTool(ctx, &CallToolParams{Name: "act"}) + if err == nil || !strings.Contains(err.Error(), "unsupported params type") { + t.Fatalf("CallTool() error = %v, want unsupported params type error", err) + } +} + +func TestMultiRoundTrip_MaxRetries(t *testing.T) { + testCases := []struct { + name string + requests InputRequestMap + wantRetries int + }{ + { + name: "load shedding", + requests: InputRequestMap{}, + wantRetries: maxLoadSheddingMultiRoundTripRetries, + }, + { + name: "input request", + requests: InputRequestMap{"confirm": &ElicitParams{Message: "Again?"}}, + wantRetries: maxMultiRoundTripRetries, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + + var serverCalls atomic.Int32 + srv := NewServer(testImpl, nil) + AddTool(srv, &Tool{Name: "loop"}, func(ctx context.Context, req *CallToolRequest, input struct{}) (*CallToolResult, any, error) { + serverCalls.Add(1) + return &CallToolResult{InputRequests: tc.requests, RequestState: "loop-state"}, nil, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + ElicitationHandler: func(_ context.Context, req *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + + _, err := conn.CallTool(ctx, &CallToolParams{Name: "loop"}) + if err == nil { + t.Fatal("CallTool() err = nil, want error for exceeded max retries") + } + if serverCalls.Load() != int32(tc.wantRetries) { + t.Errorf("serverCalls = %d, want %d", serverCalls.Load(), tc.wantRetries) + } + }) + } +} + +func TestMultiRoundTrip_ServerMiddleware(t *testing.T) { + // multiRoundTripToolHandler returns a ToolHandler (plain, non-generic) that requests + // the given inputRequests on the first call and returns the fulfilled + // response IDs on the second. + multiRoundTripToolHandler := func(inputRequests InputRequestMap) ToolHandler { + return func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + if len(req.Params.InputResponses) == 0 { + return &CallToolResult{ + InputRequests: inputRequests, + RequestState: "state-1", + }, nil + } + var ids []string + for id := range req.Params.InputResponses { + ids = append(ids, id) + } + slices.Sort(ids) + content := &TextContent{Text: fmt.Sprintf("%v", ids)} + return &CallToolResult{Content: []Content{content}}, nil + } + } + + tests := []struct { + name string + inputRequests InputRequestMap + wantText string + }{ + { + name: "elicit via ToolHandler", + inputRequests: InputRequestMap{ + "confirm": &ElicitParams{Message: "Sure?"}, + }, + wantText: "[confirm]", + }, + { + name: "elicit and listRoots via ToolHandler", + inputRequests: InputRequestMap{ + "confirm": &ElicitParams{Message: "OK?"}, + "roots": &ListRootsParams{}, + }, + wantText: "[confirm roots]", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddTool( + &Tool{Name: "act", InputSchema: &jsonschema.Schema{Type: "object"}}, + multiRoundTripToolHandler(tt.inputRequests), + ) + + // Connect with an OLD protocol version where multi-round-trip is not supported. + // The server middleware should handle it transparently. + st, ct := NewInMemoryTransports() + ss, err := srv.Connect(t.Context(), st, nil) + if err != nil { + t.Fatalf("server.Connect() error = %v", err) + } + t.Cleanup(func() { _ = ss.Close() }) + + c := NewClient(testImpl, &ClientOptions{ + MultiRoundTrip: &MultiRoundTripOptions{Disabled: true}, + ElicitationHandler: func(_ context.Context, req *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + c.AddRoots(&Root{URI: "file:///workspace", Name: "workspace"}) + cs, err := c.Connect(t.Context(), ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("client.Connect() error = %v", err) + } + t.Cleanup(func() { _ = cs.Close() }) + + res, err := cs.CallTool(ctx, &CallToolParams{Name: "act"}) + if err != nil { + t.Fatalf("CallTool() error = %v", err) + } + if got := res.Content[0].(*TextContent).Text; got != tt.wantText { + t.Errorf("result text = %q, want %q", got, tt.wantText) + } + }) + } +} + +// TestMultiRoundTrip_ServerMiddleware_GetPrompt verifies the server-side +// shim retry for prompts: the rebuilt retry request must keep its concrete +// *ServerRequest[*GetPromptParams] type for the receiving dispatch. +func TestMultiRoundTrip_ServerMiddleware_GetPrompt(t *testing.T) { + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddPrompt(&Prompt{Name: "review"}, func(_ context.Context, req *GetPromptRequest) (*GetPromptResult, error) { + if len(req.Params.InputResponses) == 0 { + return &GetPromptResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Sure?"}}, + RequestState: "prompt-state", + }, nil + } + return &GetPromptResult{ + Messages: []*PromptMessage{{Role: "user", Content: &TextContent{Text: "approved"}}}, + }, nil + }) + + conn := mustConnectOldProtocol(t, srv, &ClientOptions{ + MultiRoundTrip: &MultiRoundTripOptions{Disabled: true}, + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + + res, err := conn.GetPrompt(ctx, &GetPromptParams{Name: "review"}) + if err != nil { + t.Fatalf("GetPrompt() error = %v", err) + } + if got := res.Messages[0].Content.(*TextContent).Text; got != "approved" { + t.Errorf("message text = %q, want %q", got, "approved") + } +} + +func TestMultiRoundTrip_GetPrompt_AutoRetry(t *testing.T) { + + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddPrompt(&Prompt{Name: "review"}, func(_ context.Context, req *GetPromptRequest) (*GetPromptResult, error) { + if len(req.Params.InputResponses) == 0 { + return &GetPromptResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Include sensitive data?"}}, + RequestState: "prompt-state", + }, nil + } + return &GetPromptResult{ + Description: "Code review prompt", + Messages: []*PromptMessage{{Role: "user", Content: &TextContent{Text: "review this code"}}}, + }, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + + res, err := conn.GetPrompt(ctx, &GetPromptParams{Name: "review"}) + if err != nil { + t.Fatalf("GetPrompt() error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after auto-retry, want false") + } + if len(res.Messages) != 1 { + t.Fatalf("len(res.Messages) = %d, want 1", len(res.Messages)) + } + if got := res.Messages[0].Content.(*TextContent).Text; got != "review this code" { + t.Errorf("message text = %q, want %q", got, "review this code") + } +} + +func TestMultiRoundTrip_GetPrompt_ManualRetry(t *testing.T) { + + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddPrompt(&Prompt{Name: "review"}, func(_ context.Context, req *GetPromptRequest) (*GetPromptResult, error) { + if len(req.Params.InputResponses) == 0 { + return &GetPromptResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Include sensitive data?"}}, + RequestState: "prompt-state", + }, nil + } + return &GetPromptResult{ + Description: "Code review prompt", + Messages: []*PromptMessage{{Role: "user", Content: &TextContent{Text: "review this code"}}}, + }, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + MultiRoundTrip: &MultiRoundTripOptions{Disabled: true}, + }) + + res, err := conn.GetPrompt(ctx, &GetPromptParams{Name: "review"}) + if err != nil { + t.Fatalf("GetPrompt() error = %v", err) + } + if !res.NeedsInput() { + t.Fatal("NeedsInput() = false, want true") + } + if _, ok := res.InputRequests["confirm"].(*ElicitParams); !ok { + t.Fatalf("InputRequests[confirm] type = %T, want *ElicitParams", res.InputRequests["confirm"]) + } + + res, err = conn.GetPrompt(ctx, &GetPromptParams{ + Name: "review", + InputResponses: InputResponseMap{"confirm": &ElicitResult{Action: "accept"}}, + RequestState: res.RequestState, + }) + if err != nil { + t.Fatalf("GetPrompt() follow-up error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after follow-up, want false") + } + if len(res.Messages) != 1 { + t.Fatalf("len(res.Messages) = %d, want 1", len(res.Messages)) + } +} + +func TestMultiRoundTrip_ReadResource_AutoRetry(t *testing.T) { + + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddResource(&Resource{URI: "test://data", Name: "data"}, func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + if len(req.Params.InputResponses) == 0 { + return &ReadResourceResult{ + InputRequests: InputRequestMap{"auth": &ElicitParams{Message: "Authenticate?"}}, + RequestState: "resource-state", + }, nil + } + return &ReadResourceResult{ + Contents: []*ResourceContents{{URI: "test://data", Text: "resource data"}}, + }, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "accept"}, nil + }, + }) + + res, err := conn.ReadResource(ctx, &ReadResourceParams{URI: "test://data"}) + if err != nil { + t.Fatalf("ReadResource() error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after auto-retry, want false") + } + if len(res.Contents) != 1 { + t.Fatalf("len(res.Contents) = %d, want 1", len(res.Contents)) + } + if got := res.Contents[0].Text; got != "resource data" { + t.Errorf("resource text = %q, want %q", got, "resource data") + } +} + +func TestMultiRoundTrip_ReadResource_ManualRetry(t *testing.T) { + + ctx := context.Background() + + srv := NewServer(testImpl, nil) + srv.AddResource(&Resource{URI: "test://data", Name: "data"}, func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + if len(req.Params.InputResponses) == 0 { + return &ReadResourceResult{ + InputRequests: InputRequestMap{"auth": &ElicitParams{Message: "Authenticate?"}}, + RequestState: "resource-state", + }, nil + } + return &ReadResourceResult{ + Contents: []*ResourceContents{{URI: "test://data", Text: "resource data"}}, + }, nil + }) + + conn := mustConnect(t, srv, &ClientOptions{ + MultiRoundTrip: &MultiRoundTripOptions{Disabled: true}, + }) + + res, err := conn.ReadResource(ctx, &ReadResourceParams{URI: "test://data"}) + if err != nil { + t.Fatalf("ReadResource() error = %v", err) + } + if !res.NeedsInput() { + t.Fatal("NeedsInput() = false, want true") + } + if _, ok := res.InputRequests["auth"].(*ElicitParams); !ok { + t.Fatalf("InputRequests[auth] type = %T, want *ElicitParams", res.InputRequests["auth"]) + } + + res, err = conn.ReadResource(ctx, &ReadResourceParams{ + URI: "test://data", + InputResponses: InputResponseMap{"auth": &ElicitResult{Action: "accept"}}, + RequestState: res.RequestState, + }) + if err != nil { + t.Fatalf("ReadResource() follow-up error = %v", err) + } + if res.NeedsInput() { + t.Fatal("NeedsInput() = true after follow-up, want false") + } + if len(res.Contents) != 1 { + t.Fatalf("len(res.Contents) = %d, want 1", len(res.Contents)) + } +} + +// TestSetMultiRoundTripRetryParams pins the properties the retry-request +// builder must uphold: the original params are never mutated, the copy +// carries the responses and state, the server-side rebuild preserves the +// request's concrete *ServerRequest[P] type (the receiving dispatch +// type-asserts on it) and its Extra, and an unrecognized params type is an +// explicit error. The server readResource case matters here because the +// end-to-end shim path cannot reach it (the read handler's nil-contents +// validation rejects input-required results for old-protocol clients). +func TestSetMultiRoundTripRetryParams(t *testing.T) { + ss := &ServerSession{server: NewServer(testImpl, nil)} + extra := &RequestExtra{} + responses := InputResponseMap{"confirm": &ElicitResult{Action: "accept"}} + + // checkRetry asserts the returned params carry the retry fields while + // the original params were left untouched. + checkRetry := func(t *testing.T, orig, retry Params) { + t.Helper() + switch p := retry.(type) { + case *CallToolParams: + if len(p.InputResponses) != 1 || p.RequestState != "s1" { + t.Errorf("retry params = %+v, want responses and state carried", p) + } + case *CallToolParamsRaw: + if len(p.InputResponses) != 1 || p.RequestState != "s1" { + t.Errorf("retry params = %+v, want responses and state carried", p) + } + case *GetPromptParams: + if len(p.InputResponses) != 1 || p.RequestState != "s1" { + t.Errorf("retry params = %+v, want responses and state carried", p) + } + case *ReadResourceParams: + if len(p.InputResponses) != 1 || p.RequestState != "s1" { + t.Errorf("retry params = %+v, want responses and state carried", p) + } + default: + t.Errorf("retry params type = %T, want one of the multi-round-trip params types", p) + } + switch p := orig.(type) { + case *CallToolParams: + if p.InputResponses != nil || p.RequestState != "" { + t.Errorf("original params mutated: %+v", p) + } + case *CallToolParamsRaw: + if p.InputResponses != nil || p.RequestState != "" { + t.Errorf("original params mutated: %+v", p) + } + case *GetPromptParams: + if p.InputResponses != nil || p.RequestState != "" { + t.Errorf("original params mutated: %+v", p) + } + case *ReadResourceParams: + if p.InputResponses != nil || p.RequestState != "" { + t.Errorf("original params mutated: %+v", p) + } + } + } + + t.Run("server callTool", func(t *testing.T) { + orig := &CallToolParamsRaw{Name: "act"} + req, err := setMultiRoundTripRetryParams(methodCallTool, &ServerRequest[*CallToolParamsRaw]{Session: ss, Params: orig, Extra: extra}, responses, "s1") + if err != nil { + t.Fatalf("setMultiRoundTripRetryParams() error = %v", err) + } + got, ok := req.(*ServerRequest[*CallToolParamsRaw]) + if !ok { + t.Fatalf("retry request type = %T, want *ServerRequest[*CallToolParamsRaw]", req) + } + if got.Extra != extra { + t.Error("retry request lost Extra") + } + checkRetry(t, orig, got.Params) + }) + + t.Run("server getPrompt", func(t *testing.T) { + orig := &GetPromptParams{Name: "review"} + req, err := setMultiRoundTripRetryParams(methodGetPrompt, &ServerRequest[*GetPromptParams]{Session: ss, Params: orig, Extra: extra}, responses, "s1") + if err != nil { + t.Fatalf("setMultiRoundTripRetryParams() error = %v", err) + } + got, ok := req.(*ServerRequest[*GetPromptParams]) + if !ok { + t.Fatalf("retry request type = %T, want *ServerRequest[*GetPromptParams]", req) + } + if got.Extra != extra { + t.Error("retry request lost Extra") + } + checkRetry(t, orig, got.Params) + }) + + t.Run("server readResource", func(t *testing.T) { + orig := &ReadResourceParams{URI: "test://data"} + req, err := setMultiRoundTripRetryParams(methodReadResource, &ServerRequest[*ReadResourceParams]{Session: ss, Params: orig, Extra: extra}, responses, "s1") + if err != nil { + t.Fatalf("setMultiRoundTripRetryParams() error = %v", err) + } + got, ok := req.(*ServerRequest[*ReadResourceParams]) + if !ok { + t.Fatalf("retry request type = %T, want *ServerRequest[*ReadResourceParams]", req) + } + if got.Extra != extra { + t.Error("retry request lost Extra") + } + checkRetry(t, orig, got.Params) + }) + + t.Run("client params", func(t *testing.T) { + for _, orig := range []Params{ + &CallToolParams{Name: "act"}, + &GetPromptParams{Name: "review"}, + &ReadResourceParams{URI: "test://data"}, + } { + req, err := setMultiRoundTripRetryParams(methodCallTool, newClientRequest[Params](nil, orig), responses, "s1") + if err != nil { + t.Fatalf("setMultiRoundTripRetryParams(%T) error = %v", orig, err) + } + got, ok := req.(*ClientRequest[Params]) + if !ok { + t.Fatalf("retry request type = %T, want *ClientRequest[Params]", req) + } + checkRetry(t, orig, got.Params) + } + }) + + t.Run("unsupported params type", func(t *testing.T) { + wrapped := wrappedCallToolParams{&CallToolParams{Name: "act"}} + _, err := setMultiRoundTripRetryParams(methodCallTool, newClientRequest[Params](nil, wrapped), responses, "s1") + if err == nil || !strings.Contains(err.Error(), "unsupported params type") { + t.Fatalf("setMultiRoundTripRetryParams() error = %v, want unsupported params type error", err) + } + }) +} + +func mustConnect(t *testing.T, s *Server, clientOpts *ClientOptions) *ClientSession { + t.Helper() + + st, ct := NewInMemoryTransports() + ss, err := s.Connect(t.Context(), st, nil) + if err != nil { + t.Fatalf("server.Connect() error = %v", err) + } + t.Cleanup(func() { + _ = ss.Close() + }) + + c := NewClient(testImpl, clientOpts) + cs, err := c.Connect(t.Context(), ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client.Connect() error = %v", err) + } + t.Cleanup(func() { + _ = cs.Close() + }) + return cs +} + +// mustConnectOldProtocol is mustConnect for a client on a protocol version +// that predates multi-round-trip: the server-side shim, not the client +// middleware, fulfills input requests. +func mustConnectOldProtocol(t *testing.T, s *Server, clientOpts *ClientOptions) *ClientSession { + t.Helper() + + st, ct := NewInMemoryTransports() + ss, err := s.Connect(t.Context(), st, nil) + if err != nil { + t.Fatalf("server.Connect() error = %v", err) + } + t.Cleanup(func() { + _ = ss.Close() + }) + + c := NewClient(testImpl, clientOpts) + cs, err := c.Connect(t.Context(), ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("client.Connect() error = %v", err) + } + t.Cleanup(func() { + _ = cs.Close() + }) + return cs +} diff --git a/mcp/protocol.go b/mcp/protocol.go index 1312dfb..c340e60 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -4,17 +4,205 @@ package mcp -// Protocol types for version 2025-06-18. -// To see the schema changes from the previous version, run: -// -// prefix=https://raw.githubusercontent.com/modelcontextprotocol/modelcontextprotocol/refs/heads/main/schema -// sdiff -l <(curl $prefix/2025-03-26/schema.ts) <(curl $prefix/2025/06-18/schema.ts) - import ( "encoding/json" "fmt" + "maps" + + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" + "github.com/modelcontextprotocol/go-sdk/internal/mcpgodebug" ) +// resultType indicates whether a result is complete or requires further input +// from the client via the multi round-trip request protocol. +type resultType string + +const ( + // resultTypeComplete indicates the result is final. + // This is the default when resultType is empty. + resultTypeComplete resultType = "complete" + + // resultTypeInputRequired indicates the server needs additional client + // input before it can complete the request. The client should fulfill the + // InputRequests and retry the call with the responses. + resultTypeInputRequired resultType = "input_required" +) + +type completeResultWithType struct { + ResultType resultType `json:"resultType,omitempty"` +} + +func (r *completeResultWithType) setResultType(rt resultType) { r.ResultType = rt } +func (*completeResultWithType) isCompleteResult() {} + +type completeResultResponse interface { + setResultType(resultType) + isCompleteResult() +} + +// annotateResultType sets the resultType that protocol 2026-07-28 requires on +// every result. +func annotateResultType(res Result) { + switch r := res.(type) { + case completeResultResponse: + r.setResultType(resultTypeComplete) + case multiRoundTripResponse: + // These results are complete or input_required, so label them by + // whether the handler asked for more client input. + if r.inputRequests() != nil { + r.setResultType(resultTypeInputRequired) + } else { + r.setResultType(resultTypeComplete) + } + } +} + +// InputRequest is a type for parameters that a server can include in the response +// to request input from client (SEP-2322). Implementations are [*ElicitParams], +// [*CreateMessageParams], and [*ListRootsParams]. +type InputRequest interface{ isInputRequest() } + +// InputRequestMap maps server-assigned request IDs to [InputRequest] values. +// It is used in result types to tell the client what input the server needs. +type InputRequestMap map[string]InputRequest + +func (m InputRequestMap) MarshalJSON() ([]byte, error) { + if m == nil { + return json.Marshal(map[string]any(nil)) + } + type wire struct { + Method string `json:"method"` + Params InputRequest `json:"params,omitempty"` + } + typeToMethod := func(v InputRequest) (string, error) { + switch v.(type) { + case *ElicitParams: + return methodElicit, nil + case *CreateMessageParams, *CreateMessageWithToolsParams: + return methodCreateMessage, nil + case *ListRootsParams: + return methodListRoots, nil + default: + return "", fmt.Errorf("unsupported type: %T", v) + } + } + converted := map[string]*wire{} + for k, v := range m { + method, err := typeToMethod(v) + if err != nil { + return nil, err + } + if ep, ok := v.(*ElicitParams); ok { + v = ep.inferElicitMode() + } + converted[k] = &wire{Method: method, Params: v} + } + return json.Marshal(converted) +} + +func (m *InputRequestMap) UnmarshalJSON(data []byte) error { + type raw struct { + Method string `json:"method"` + Params json.RawMessage `json:"params"` + } + var rawMap map[string]*raw + if err := json.Unmarshal(data, &rawMap); err != nil { + return err + } + if rawMap == nil { + return nil + } + result := make(InputRequestMap, len(rawMap)) + for k, raw := range rawMap { + switch raw.Method { + case methodElicit: + var p ElicitParams + if err := json.Unmarshal(raw.Params, &p); err != nil { + return err + } + result[k] = &p + case methodCreateMessage: + var p CreateMessageWithToolsParams + if err := json.Unmarshal(raw.Params, &p); err != nil { + return err + } + result[k] = &p + case methodListRoots: + var p ListRootsParams + if err := json.Unmarshal(raw.Params, &p); err != nil { + return err + } + result[k] = &p + default: + return fmt.Errorf("unsupported InputRequest method: %q", raw.Method) + } + } + *m = result + return nil +} + +// InputResponse is a type for results that a client sends back when fulfilling +// a server input request (SEP-2322). Implementations are [*ElicitResult], +// [*CreateMessageResult], and [*ListRootsResult]. +type InputResponse interface{ isInputResponse() } + +// InputResponseMap maps request IDs (from [InputRequestMap]) to [InputResponse] +// values. It is used in params types when retrying a call after an +// input-required result. +type InputResponseMap map[string]InputResponse + +func (m *InputResponseMap) UnmarshalJSON(data []byte) error { + var rawMap map[string]json.RawMessage + if err := json.Unmarshal(data, &rawMap); err != nil { + return err + } + result := make(InputResponseMap, len(rawMap)) + for k, raw := range rawMap { + v, err := unmarshalInputResponse(raw) + if err != nil { + return fmt.Errorf("inputResponses[%q]: %w", k, err) + } + result[k] = v + } + *m = result + return nil +} + +// unmarshalInputResponse determines the concrete InputResponse type from the +// JSON structure by searching for a discriminating key in a raw message. +func unmarshalInputResponse(data json.RawMessage) (InputResponse, error) { + var probe struct { + Action json.RawMessage `json:"action"` + Role json.RawMessage `json:"role"` + Roots json.RawMessage `json:"roots"` + } + if err := json.Unmarshal(data, &probe); err != nil { + return nil, err + } + switch { + case probe.Roots != nil: + var p ListRootsResult + if err := json.Unmarshal(data, &p); err != nil { + return nil, err + } + return &p, nil + case probe.Action != nil: + var p ElicitResult + if err := json.Unmarshal(data, &p); err != nil { + return nil, err + } + return &p, nil + case probe.Role != nil: + var p CreateMessageWithToolsResult + if err := json.Unmarshal(data, &p); err != nil { + return nil, err + } + return &p, nil + default: + return nil, fmt.Errorf(`cannot determine InputResponse type: expected "action", "role", or "roots" key`) + } +} + // Optional annotations for the client. The client can use annotations to inform // how objects are used or displayed. type Annotations struct { @@ -48,6 +236,13 @@ type CallToolParams struct { // Arguments holds the tool arguments. It can hold any value that can be // marshaled to JSON. Arguments any `json:"arguments,omitempty"` + + // InputResponses maps input request IDs to responses. Set it when the + // client retries a call that asked for input. See [CallToolResult.NeedsInput]. + InputResponses InputResponseMap `json:"inputResponses,omitempty"` + // RequestState is the opaque state from the previous input-required result. + // The client must echo this back when retrying. + RequestState string `json:"requestState,omitempty"` } // CallToolParamsRaw is passed to tool handlers on the server. Its arguments @@ -63,6 +258,13 @@ type CallToolParamsRaw struct { // is the responsibility of the tool handler to unmarshal and validate the // Arguments (see [AddTool]). Arguments json.RawMessage `json:"arguments,omitempty"` + + // InputResponses maps input request IDs to responses. Set it when the + // client retries a call that asked for input. See [CallToolResult.NeedsInput]. + InputResponses InputResponseMap `json:"inputResponses,omitempty"` + // RequestState is the opaque state from the previous input-required result. + // The client must echo this back when retrying. + RequestState string `json:"requestState,omitempty"` } // A CallToolResult is the server's response to a tool call. @@ -84,7 +286,11 @@ type CallToolResult struct { Content []Content `json:"content"` // StructuredContent is an optional value that represents the structured - // result of the tool call. It must marshal to a JSON object. + // result of the tool call. Per SEP-2106, it may marshal to any valid JSON + // value (object, array, or primitive) conforming to the tool's + // [Tool.OutputSchema]. + // Numbers received from the wire are represented as [json.Number] to preserve + // their exact values. // // When using a [ToolHandlerFor] with structured output, you should not // populate this field. It will be automatically populated with the typed Out @@ -109,6 +315,24 @@ type CallToolResult struct { // the Content field. IsError bool `json:"isError,omitempty"` + // InputRequests is a map of server-assigned IDs to input requests. + // It is set only when the result needs more client input. + // The client must fulfill these and echo the IDs back in InputResponses + // when it retries the call. + InputRequests InputRequestMap `json:"inputRequests,omitempty"` + + // RequestState is an opaque string the client must echo back when + // retrying after an input-required result. Servers use this to carry + // context between independent requests. + // + // Unauthenticated servers must encrypt, sign and verify this value. + RequestState string `json:"requestState,omitempty"` + + // resultType records whether the call finished or needs more client input. + // Empty or resultTypeComplete means it finished. resultTypeInputRequired + // means the client must fulfill InputRequests and retry. + // See [CallToolResult.NeedsInput]. + resultType resultType // The error passed to setError, if any. // It is not marshaled, and therefore it is only visible on the server. // Its only use is in server sending middleware, where it can be accessed @@ -116,46 +340,107 @@ type CallToolResult struct { err error } -// TODO(#64): consider exposing setError (and getError), by adding an error -// field on CallToolResult. -func (r *CallToolResult) setError(err error) { - r.Content = []Content{&TextContent{Text: err.Error()}} +// SetError sets the error for the tool result and sets IsError to true. +// If Content has not already been populated, it is set to the error text. +// If Content has already been populated, it is left unchanged, allowing callers +// to provide a user-friendly message while still recording the underlying error +// for inspection via [GetError] in server middleware. +func (r *CallToolResult) SetError(err error) { + if len(r.Content) == 0 { + r.Content = []Content{&TextContent{Text: err.Error()}} + } r.IsError = true r.err = err } -// getError returns the error set with setError, or nil if none. +// GetError returns the error set with SetError, or nil if none. // This function always returns nil on clients. -func (r *CallToolResult) getError() error { +func (r *CallToolResult) GetError() error { return r.err } -func (*CallToolResult) isResult() {} +func (*CallToolResult) isResult() {} +func (x *CallToolResult) isNil() bool { return x == nil } + +func (r *CallToolResult) setResultType(rt resultType) { r.resultType = rt } +func (r *CallToolResult) requestState() string { return r.RequestState } +func (r *CallToolResult) inputRequests() map[string]InputRequest { + if r == nil { + return nil + } + return r.InputRequests +} +func (r *CallToolResult) hasContent() bool { + return len(r.Content) > 0 || r.StructuredContent != nil +} + +// NeedsInput reports whether this result requires further client input. +// This is true when the server returned a resultType of "input_required". +// When NeedsInput returns true, check InputRequests for the set of +// requests the server needs fulfilled before retrying the call. +// An empty InputRequests with NeedsInput true indicates load-shedding. +func (r *CallToolResult) NeedsInput() bool { return r.resultType == resultTypeInputRequired } + +// structuredcontentfloat64 is a compatibility parameter that restores the +// previous behavior of decoding numbers in [CallToolResult.StructuredContent] +// as float64 values. By default, numbers are decoded as [json.Number] to avoid +// losing precision. The option will be removed in the 1.11.0 version of the SDK. +var structuredcontentfloat64 = mcpgodebug.Value("structuredcontentfloat64") + +func (x *CallToolResult) MarshalJSON() ([]byte, error) { + type res CallToolResult // avoid recursion + type wire struct { + res + ResultType resultType `json:"resultType,omitempty"` + InputRequests json.RawMessage `json:"inputRequests,omitempty"` // shadows res.InputRequests + } + w := wire{res: res(*x), ResultType: x.resultType} + if x.InputRequests != nil { + ir, err := json.Marshal(x.InputRequests) + if err != nil { + return nil, err + } + w.InputRequests = ir + } + return json.Marshal(w) +} -// UnmarshalJSON handles the unmarshalling of content into the Content -// interface. func (x *CallToolResult) UnmarshalJSON(data []byte) error { type res CallToolResult // avoid recursion var wire struct { res - Content []*wireContent `json:"content"` + Content []*wireContent `json:"content"` + StructuredContent json.RawMessage `json:"structuredContent"` + ResultType resultType `json:"resultType"` } - if err := json.Unmarshal(data, &wire); err != nil { + if err := internaljson.Unmarshal(data, &wire); err != nil { return err } + if len(wire.StructuredContent) > 0 { + unmarshal := internaljson.UnmarshalUseNumber + if structuredcontentfloat64 == "1" { + unmarshal = internaljson.Unmarshal + } + if err := unmarshal(wire.StructuredContent, &wire.res.StructuredContent); err != nil { + return err + } + } var err error if wire.res.Content, err = contentsFromWire(wire.Content, nil); err != nil { return err } + wire.res.resultType = wire.ResultType *x = CallToolResult(wire.res) return nil } func (x *CallToolParams) isParams() {} +func (x *CallToolParams) isNil() bool { return x == nil } func (x *CallToolParams) GetProgressToken() any { return getProgressToken(x) } func (x *CallToolParams) SetProgressToken(t any) { setProgressToken(x, t) } func (x *CallToolParamsRaw) isParams() {} +func (x *CallToolParamsRaw) isNil() bool { return x == nil } func (x *CallToolParamsRaw) GetProgressToken() any { return getProgressToken(x) } func (x *CallToolParamsRaw) SetProgressToken(t any) { setProgressToken(x, t) } @@ -174,26 +459,139 @@ type CancelledParams struct { } func (x *CancelledParams) isParams() {} +func (x *CancelledParams) isNil() bool { return x == nil } func (x *CancelledParams) GetProgressToken() any { return getProgressToken(x) } func (x *CancelledParams) SetProgressToken(t any) { setProgressToken(x, t) } +// RootCapabilities describes a client's support for roots. +// +// Deprecated: the roots feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +type RootCapabilities struct { + // ListChanged reports whether the client supports notifications for + // changes to the roots list. + ListChanged bool `json:"listChanged,omitempty"` +} + // Capabilities a client may support. Known capabilities are defined here, in // this schema, but this is not a closed set: any client can define its own, // additional capabilities. type ClientCapabilities struct { - // Experimental, non-standard capabilities that the client supports. + // NOTE: any addition to ClientCapabilities must also be reflected in + // [ClientCapabilities.clone]. + + // Experimental reports non-standard capabilities that the client supports. + // The caller should not modify the map after assigning it. Experimental map[string]any `json:"experimental,omitempty"` - // Present if the client supports listing roots. + // Extensions reports extensions that the client supports. + // Keys are extension identifiers in "{vendor-prefix}/{extension-name}" format. + // Values are per-extension settings objects; use [ClientCapabilities.AddExtension] + // to ensure nil settings are normalized to empty objects. + // The caller should not modify the map or its values after assigning it. + Extensions map[string]any `json:"extensions,omitempty"` + // Roots describes the client's support for roots. + // + // Deprecated: use RootsV2. As described in #607, Roots should have been a + // pointer to a RootCapabilities value. Roots will be continue to be + // populated, but any new fields will only be added in the RootsV2 field. + // + // The roots feature itself is also deprecated by SEP-2577; see RootsV2. Roots struct { - // Whether the client supports notifications for changes to the roots list. + // ListChanged reports whether the client supports notifications for + // changes to the roots list. ListChanged bool `json:"listChanged,omitempty"` } `json:"roots,omitempty"` - // Present if the client supports sampling from an LLM. + // RootsV2 is present if the client supports roots. When capabilities are + // explicitly configured via [ClientOptions.Capabilities]. + // + // Deprecated: the roots feature is deprecated as of protocol version + // 2026-07-28 (SEP-2577). It remains functional during the deprecation + // window (at least twelve months). See + // https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. + RootsV2 *RootCapabilities `json:"-"` + // Sampling is present if the client supports sampling from an LLM. + // + // Deprecated: the sampling feature is deprecated as of protocol version + // 2026-07-28 (SEP-2577). It remains functional during the deprecation + // window (at least twelve months). See + // https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. Sampling *SamplingCapabilities `json:"sampling,omitempty"` - // Present if the client supports elicitation from the server. + // Elicitation is present if the client supports elicitation from the server. Elicitation *ElicitationCapabilities `json:"elicitation,omitempty"` } +// AddExtension adds an extension with the given name and settings. +// If settings is nil, an empty map is used to ensure valid JSON serialization +// (the spec requires an object, not null). +// The settings map should not be modified after the call. +func (c *ClientCapabilities) AddExtension(name string, settings map[string]any) { + if c.Extensions == nil { + c.Extensions = make(map[string]any) + } + if settings == nil { + settings = map[string]any{} + } + c.Extensions[name] = settings +} + +// clone returns a copy of the ClientCapabilities. +// Values in the Extensions and Experimental maps are shallow-copied. +func (c *ClientCapabilities) clone() *ClientCapabilities { + cp := *c + cp.Experimental = maps.Clone(c.Experimental) + cp.Extensions = maps.Clone(c.Extensions) + cp.RootsV2 = shallowClone(c.RootsV2) + if c.Sampling != nil { + x := *c.Sampling + x.Tools = shallowClone(c.Sampling.Tools) + x.Context = shallowClone(c.Sampling.Context) + cp.Sampling = &x + } + if c.Elicitation != nil { + x := *c.Elicitation + x.Form = shallowClone(c.Elicitation.Form) + x.URL = shallowClone(c.Elicitation.URL) + cp.Elicitation = &x + } + return &cp +} + +// shallowClone returns a shallow clone of *p, or nil if p is nil. +func shallowClone[T any](p *T) *T { + if p == nil { + return nil + } + x := *p + return &x +} + +func (c *ClientCapabilities) toV2() *clientCapabilitiesV2 { + return &clientCapabilitiesV2{ + ClientCapabilities: *c, + Roots: c.RootsV2, + } +} + +// clientCapabilitiesV2 is a version of ClientCapabilities that fixes the bug +// described in #607: Roots should have been a pointer to value type +// RootCapabilities. +type clientCapabilitiesV2 struct { + ClientCapabilities + Roots *RootCapabilities `json:"roots,omitempty"` +} + +func (c *clientCapabilitiesV2) toV1() *ClientCapabilities { + caps := c.ClientCapabilities + caps.RootsV2 = c.Roots + // Sync Roots from RootsV2 for backward compatibility (#607). + if caps.RootsV2 != nil { + caps.Roots = *caps.RootsV2 + } + return &caps +} + type CompleteParamsArgument struct { // The name of the argument Name string `json:"name"` @@ -220,7 +618,7 @@ type CompleteReference struct { func (r *CompleteReference) UnmarshalJSON(data []byte) error { type wireCompleteReference CompleteReference // for naive unmarshaling var r2 wireCompleteReference - if err := json.Unmarshal(data, &r2); err != nil { + if err := internaljson.Unmarshal(data, &r2); err != nil { return err } switch r2.Type { @@ -267,7 +665,8 @@ type CompleteParams struct { Ref *CompleteReference `json:"ref"` } -func (*CompleteParams) isParams() {} +func (x *CompleteParams) isParams() {} +func (x *CompleteParams) isNil() bool { return x == nil } type CompletionResultDetails struct { HasMore bool `json:"hasMore,omitempty"` @@ -277,20 +676,34 @@ type CompletionResultDetails struct { // The server's response to a completion/complete request type CompleteResult struct { + completeResultWithType // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. Meta `json:"_meta,omitempty"` Completion CompletionResultDetails `json:"completion"` } -func (*CompleteResult) isResult() {} +func (*CompleteResult) isResult() {} +func (x *CompleteResult) isNil() bool { return x == nil } +// CreateMessageParams holds parameters for a sampling/createMessage request. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type CreateMessageParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. Meta `json:"_meta,omitempty"` // A request to include context from one or more MCP servers (including the // caller), to be attached to the prompt. The client may ignore this request. + // + // The default is "none". The values "thisServer" and "allServers" are + // deprecated as of protocol version 2025-11-25 (SEP-2596) and will be + // removed no later than the sampling feature itself (SEP-2577). Servers + // SHOULD omit this field or use "none". See + // https://modelcontextprotocol.io/seps/2596-feature-lifecycle-and-deprecation-policy. IncludeContext string `json:"includeContext,omitempty"` // The maximum number of tokens to sample, as requested by the server. The // client may choose to sample fewer tokens than requested. @@ -310,13 +723,132 @@ type CreateMessageParams struct { } func (x *CreateMessageParams) isParams() {} +func (x *CreateMessageParams) isInputRequest() {} +func (x *CreateMessageParams) isNil() bool { return x == nil } func (x *CreateMessageParams) GetProgressToken() any { return getProgressToken(x) } func (x *CreateMessageParams) SetProgressToken(t any) { setProgressToken(x, t) } +// CreateMessageWithToolsParams is a sampling request that includes tools. +// It extends the basic [CreateMessageParams] fields with tools, tool choice, +// and messages that support array content (for parallel tool calls). +// +// Use with [ServerSession.CreateMessageWithTools]. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +type CreateMessageWithToolsParams struct { + Meta `json:"_meta,omitempty"` + IncludeContext string `json:"includeContext,omitempty"` + MaxTokens int64 `json:"maxTokens"` + // Messages supports array content for tool_use and tool_result blocks. + Messages []*SamplingMessageV2 `json:"messages"` + Metadata any `json:"metadata,omitempty"` + ModelPreferences *ModelPreferences `json:"modelPreferences,omitempty"` + StopSequences []string `json:"stopSequences,omitempty"` + SystemPrompt string `json:"systemPrompt,omitempty"` + Temperature float64 `json:"temperature,omitempty"` + // Tools is the list of tools available for the model to use. + Tools []*Tool `json:"tools,omitempty"` + // ToolChoice controls how the model should use tools. + ToolChoice *ToolChoice `json:"toolChoice,omitempty"` +} + +func (x *CreateMessageWithToolsParams) isParams() {} +func (x *CreateMessageWithToolsParams) isInputRequest() {} +func (x *CreateMessageWithToolsParams) isNil() bool { return x == nil } +func (x *CreateMessageWithToolsParams) GetProgressToken() any { return getProgressToken(x) } +func (x *CreateMessageWithToolsParams) SetProgressToken(t any) { setProgressToken(x, t) } + +// toBase converts to CreateMessageParams by taking the content block from each +// message. Tools and ToolChoice are dropped. Returns an error if any message +// has multiple content blocks, since SamplingMessage only supports one. +func (p *CreateMessageWithToolsParams) toBase() (*CreateMessageParams, error) { + var msgs []*SamplingMessage + for _, m := range p.Messages { + if len(m.Content) > 1 { + return nil, fmt.Errorf("message has %d content blocks; use CreateMessageWithToolsHandler to support multiple content", len(m.Content)) + } + var content Content + if len(m.Content) > 0 { + content = m.Content[0] + } + msgs = append(msgs, &SamplingMessage{Content: content, Role: m.Role}) + } + return &CreateMessageParams{ + Meta: p.Meta, + IncludeContext: p.IncludeContext, + MaxTokens: p.MaxTokens, + Messages: msgs, + Metadata: p.Metadata, + ModelPreferences: p.ModelPreferences, + StopSequences: p.StopSequences, + SystemPrompt: p.SystemPrompt, + Temperature: p.Temperature, + }, nil +} + +// SamplingMessageV2 describes a message issued to or received from an +// LLM API, supporting array content for parallel tool calls. The "V2" refers +// to the 2025-11-25 spec, which changed content from a single block to +// single-or-array. In v2 of the SDK, this will replace [SamplingMessage]. +// +// When marshaling, a single-element Content slice is marshaled as a single +// object for compatibility with pre-2025-11-25 implementations. When +// unmarshaling, a single JSON content object is accepted and wrapped in a +// one-element slice. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +type SamplingMessageV2 struct { + Content []Content `json:"content"` + Role Role `json:"role"` +} + +var samplingWithToolsAllow = map[string]bool{ + "text": true, "image": true, "audio": true, + "tool_use": true, "tool_result": true, +} + +// MarshalJSON marshals the message. A single-element Content slice is marshaled +// as a single object for backward compatibility. +func (m *SamplingMessageV2) MarshalJSON() ([]byte, error) { + if len(m.Content) == 1 { + return json.Marshal(&SamplingMessage{Content: m.Content[0], Role: m.Role}) + } + type msg SamplingMessageV2 // avoid recursion + return json.Marshal((*msg)(m)) +} + +func (m *SamplingMessageV2) UnmarshalJSON(data []byte) error { + type msg SamplingMessageV2 // avoid recursion + var wire struct { + msg + Content json.RawMessage `json:"content"` + } + if err := internaljson.Unmarshal(data, &wire); err != nil { + return err + } + var err error + if wire.msg.Content, err = unmarshalContent(wire.Content, samplingWithToolsAllow); err != nil { + return err + } + *m = SamplingMessageV2(wire.msg) + return nil +} + // The client's response to a sampling/create_message request from the server. // The client should inform the user before returning the sampled message, to // allow them to inspect the response (human in the loop) and decide whether to // allow the server to see it. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type CreateMessageResult struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. @@ -326,17 +858,25 @@ type CreateMessageResult struct { Model string `json:"model"` Role Role `json:"role"` // The reason why sampling stopped, if known. + // + // Standard values: + // - "endTurn": natural end of the assistant's turn + // - "stopSequence": a stop sequence was encountered + // - "maxTokens": reached the maximum token limit + // - "toolUse": the model wants to use one or more tools StopReason string `json:"stopReason,omitempty"` } -func (*CreateMessageResult) isResult() {} +func (*CreateMessageResult) isResult() {} +func (x *CreateMessageResult) isNil() bool { return x == nil } +func (*CreateMessageResult) isInputResponse() {} func (r *CreateMessageResult) UnmarshalJSON(data []byte) error { type result CreateMessageResult // avoid recursion var wire struct { result Content *wireContent `json:"content"` } - if err := json.Unmarshal(data, &wire); err != nil { + if err := internaljson.Unmarshal(data, &wire); err != nil { return err } var err error @@ -347,6 +887,91 @@ func (r *CreateMessageResult) UnmarshalJSON(data []byte) error { return nil } +// CreateMessageWithToolsResult is the client's response to a +// sampling/create_message request that included tools. Content is a slice to +// support parallel tool calls (multiple tool_use blocks in one response). +// +// Use [ServerSession.CreateMessageWithTools] to send a sampling request with +// tools and receive this result type. +// +// When unmarshaling, a single JSON content object is accepted and wrapped in a +// one-element slice, for compatibility with clients that return a single block. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +type CreateMessageWithToolsResult struct { + Meta `json:"_meta,omitempty"` + Content []Content `json:"content"` + Model string `json:"model"` + Role Role `json:"role"` + // The reason why sampling stopped. + // + // Standard values: "endTurn", "stopSequence", "maxTokens", "toolUse". + StopReason string `json:"stopReason,omitempty"` +} + +// createMessageWithToolsResultAllow lists content types valid in assistant responses. +// tool_result is excluded: it only appears in user messages. +var createMessageWithToolsResultAllow = map[string]bool{ + "text": true, "image": true, "audio": true, + "tool_use": true, +} + +func (*CreateMessageWithToolsResult) isResult() {} +func (x *CreateMessageWithToolsResult) isNil() bool { return x == nil } +func (*CreateMessageWithToolsResult) isInputResponse() {} + +// MarshalJSON marshals the result. When Content has a single element, it is +// marshaled as a single object for compatibility with pre-2025-11-25 +// implementations that expect a single content block. +func (r *CreateMessageWithToolsResult) MarshalJSON() ([]byte, error) { + if len(r.Content) == 1 { + return json.Marshal(&CreateMessageResult{ + Meta: r.Meta, + Content: r.Content[0], + Model: r.Model, + Role: r.Role, + StopReason: r.StopReason, + }) + } + type result CreateMessageWithToolsResult // avoid recursion + return json.Marshal((*result)(r)) +} + +func (r *CreateMessageWithToolsResult) UnmarshalJSON(data []byte) error { + type result CreateMessageWithToolsResult // avoid recursion + var wire struct { + result + Content json.RawMessage `json:"content"` + } + if err := internaljson.Unmarshal(data, &wire); err != nil { + return err + } + var err error + if wire.result.Content, err = unmarshalContent(wire.Content, createMessageWithToolsResultAllow); err != nil { + return err + } + *r = CreateMessageWithToolsResult(wire.result) + return nil +} + +// toWithTools converts a CreateMessageResult to CreateMessageWithToolsResult. +func (r *CreateMessageResult) toWithTools() *CreateMessageWithToolsResult { + var content []Content + if r.Content != nil { + content = []Content{r.Content} + } + return &CreateMessageWithToolsResult{ + Meta: r.Meta, + Content: content, + Model: r.Model, + Role: r.Role, + StopReason: r.StopReason, + } +} + type GetPromptParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. @@ -355,9 +980,16 @@ type GetPromptParams struct { Arguments map[string]string `json:"arguments,omitempty"` // The name of the prompt or prompt template. Name string `json:"name"` + + // InputResponses maps input request IDs to responses. Set it when the + // client retries a call that asked for input. See [CallToolResult.NeedsInput]. + InputResponses InputResponseMap `json:"inputResponses,omitempty"` + // RequestState is the opaque state from the previous input-required result. + RequestState string `json:"requestState,omitempty"` } func (x *GetPromptParams) isParams() {} +func (x *GetPromptParams) isNil() bool { return x == nil } func (x *GetPromptParams) GetProgressToken() any { return getProgressToken(x) } func (x *GetPromptParams) SetProgressToken(t any) { setProgressToken(x, t) } @@ -369,31 +1001,116 @@ type GetPromptResult struct { // An optional description for the prompt. Description string `json:"description,omitempty"` Messages []*PromptMessage `json:"messages"` + + // InputRequests is set when the result needs more client input. + // See [CallToolResult.InputRequests]. + InputRequests InputRequestMap `json:"inputRequests,omitempty"` + // RequestState is the opaque state for multi-round-trip retries. + // See [CallToolResult.RequestState]. + RequestState string `json:"requestState,omitempty"` + + // resultType records whether the call finished or needs more client input. + // See [CallToolResult.NeedsInput]. + resultType resultType } -func (*GetPromptResult) isResult() {} +func (*GetPromptResult) isResult() {} +func (x *GetPromptResult) isNil() bool { return x == nil } +func (r *GetPromptResult) setResultType(rt resultType) { r.resultType = rt } +func (r *GetPromptResult) requestState() string { return r.RequestState } +func (r *GetPromptResult) inputRequests() map[string]InputRequest { + if r == nil { + return nil + } + return r.InputRequests +} +func (r *GetPromptResult) hasContent() bool { return len(r.Messages) > 0 } + +// NeedsInput reports whether this result requires further client input. +// See [CallToolResult.NeedsInput] for details. +func (r *GetPromptResult) NeedsInput() bool { return r.resultType == resultTypeInputRequired } + +func (x *GetPromptResult) MarshalJSON() ([]byte, error) { + type res GetPromptResult + type wire struct { + res + ResultType resultType `json:"resultType,omitempty"` + InputRequests json.RawMessage `json:"inputRequests,omitempty"` // shadows res.InputRequests + } + w := wire{res: res(*x), ResultType: x.resultType} + if x.InputRequests != nil { + ir, err := json.Marshal(x.InputRequests) + if err != nil { + return nil, err + } + w.InputRequests = ir + } + return json.Marshal(w) +} + +func (x *GetPromptResult) UnmarshalJSON(data []byte) error { + type res GetPromptResult + var wire struct { + res + ResultType resultType `json:"resultType"` + } + if err := internaljson.Unmarshal(data, &wire); err != nil { + return err + } + wire.res.resultType = wire.ResultType + *x = GetPromptResult(wire.res) + return nil +} + +// InitializeParams is sent by the client to initialize the session. type InitializeParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. - Meta `json:"_meta,omitempty"` + Meta `json:"_meta,omitempty"` + // Capabilities describes the client's capabilities. Capabilities *ClientCapabilities `json:"capabilities"` - ClientInfo *Implementation `json:"clientInfo"` - // The latest version of the Model Context Protocol that the client supports. - // The client may decide to support older versions as well. + // ClientInfo provides information about the client. + ClientInfo *Implementation `json:"clientInfo"` + // ProtocolVersion is the latest version of the Model Context Protocol that + // the client supports. ProtocolVersion string `json:"protocolVersion"` } +func (p *InitializeParams) toV2() *initializeParamsV2 { + return &initializeParamsV2{ + InitializeParams: *p, + Capabilities: p.Capabilities.toV2(), + } +} + +// initializeParamsV2 works around the mistake in #607: Capabilities.Roots +// should have been a pointer. +type initializeParamsV2 struct { + InitializeParams + Capabilities *clientCapabilitiesV2 `json:"capabilities"` +} + +func (p *initializeParamsV2) toV1() *InitializeParams { + p1 := p.InitializeParams + if p.Capabilities != nil { + p1.Capabilities = p.Capabilities.toV1() + } + return &p1 +} + func (x *InitializeParams) isParams() {} +func (x *InitializeParams) isNil() bool { return x == nil } func (x *InitializeParams) GetProgressToken() any { return getProgressToken(x) } func (x *InitializeParams) SetProgressToken(t any) { setProgressToken(x, t) } -// After receiving an initialize request from the client, the server sends this -// response. +// InitializeResult is sent by the server in response to an initialize request +// from the client. type InitializeResult struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. - Meta `json:"_meta,omitempty"` + Meta `json:"_meta,omitempty"` + // Capabilities describes the server's capabilities. Capabilities *ServerCapabilities `json:"capabilities"` // Instructions describing how to use the server and its features. // @@ -408,15 +1125,17 @@ type InitializeResult struct { ServerInfo *Implementation `json:"serverInfo"` } -func (*InitializeResult) isResult() {} +func (*InitializeResult) isResult() {} +func (x *InitializeResult) isNil() bool { return x == nil } type InitializedParams struct { - // This property is reserved by the protocol to allow clients and servers to - // attach additional metadata to their responses. + // Meta is reserved by the protocol to allow clients and servers to attach + // additional metadata to their responses. Meta `json:"_meta,omitempty"` } func (x *InitializedParams) isParams() {} +func (x *InitializedParams) isNil() bool { return x == nil } func (x *InitializedParams) GetProgressToken() any { return getProgressToken(x) } func (x *InitializedParams) SetProgressToken(t any) { setProgressToken(x, t) } @@ -429,16 +1148,91 @@ type ListPromptsParams struct { Cursor string `json:"cursor,omitempty"` } +type DiscoverParams struct { + Meta `json:"_meta,omitempty"` +} + +func (x *DiscoverParams) isParams() {} +func (x *DiscoverParams) isNil() bool { return x == nil } +func (x *DiscoverParams) GetProgressToken() any { return getProgressToken(x) } +func (x *DiscoverParams) SetProgressToken(t any) { setProgressToken(x, t) } + +type DiscoverResult struct { + completeResultWithType + Meta `json:"_meta,omitempty"` + Cacheable + // The versions of the Model Context Protocol that the server supports. + SupportedVersions []string `json:"supportedVersions"` + // The server's capabilities. + Capabilities *ServerCapabilities `json:"capabilities"` + // Instructions describing how to use the server and its features. + Instructions string `json:"instructions,omitempty"` +} + +func (*DiscoverResult) isResult() {} +func (x *DiscoverResult) isNil() bool { return x == nil } + func (x *ListPromptsParams) isParams() {} +func (x *ListPromptsParams) isNil() bool { return x == nil } func (x *ListPromptsParams) GetProgressToken() any { return getProgressToken(x) } func (x *ListPromptsParams) SetProgressToken(t any) { setProgressToken(x, t) } func (x *ListPromptsParams) cursorPtr() *string { return &x.Cursor } +// CacheableResult is a result that supports a time-to-live (TTL) hint for +// client-side caching. +type CacheableResult interface { + Result + GetTTLMs() int + GetCacheScope() string +} + +// Cacheable describes a result that supports a time-to-live (TTL) hint for +// client-side caching. +type Cacheable struct { + // A hint from the server indicating how long (in milliseconds) the + // client MAY cache this response before re-fetching. Semantics are + // analogous to HTTP Cache-Control max-age. + // + // If 0, the response SHOULD be considered immediately stale. + // If positive, the client SHOULD consider the result fresh for this + // many milliseconds after receiving the response. + TTLMs int `json:"ttlMs"` + + // Indicates the intended scope of the cached response, analogous to + // HTTP Cache-Control: public vs Cache-Control: private. + // + // "public": Any client or intermediary MAY cache and serve the response. + // "private": Only the requesting user's client MAY cache the response. + // + // Defaults to "public" if absent. + CacheScope string `json:"cacheScope"` +} + +// GetTTLMs returns the TTL hint in milliseconds. +func (c Cacheable) GetTTLMs() int { return c.TTLMs } + +// GetCacheScope returns the cache scope. +func (c Cacheable) GetCacheScope() string { return c.CacheScope } + +// normalize fills in the protocol default for any cache field left +// unset. An absent cacheScope means "public", but the field is required on the +// wire, so the default is materialized here rather than sent empty. +// +// Values already present are preserved: this must not undo a decision made by +// a resource handler or by [ServerOptions.SetCacheable]. +func (c *Cacheable) normalize() { + if c.CacheScope == "" { + c.CacheScope = "public" + } +} + // The server's response to a prompts/list request from the client. type ListPromptsResult struct { + completeResultWithType // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. Meta `json:"_meta,omitempty"` + Cacheable // An opaque token representing the pagination position after the last returned // result. If present, there may be more results available. NextCursor string `json:"nextCursor,omitempty"` @@ -446,6 +1240,7 @@ type ListPromptsResult struct { } func (x *ListPromptsResult) isResult() {} +func (x *ListPromptsResult) isNil() bool { return x == nil } func (x *ListPromptsResult) nextCursorPtr() *string { return &x.NextCursor } type ListResourceTemplatesParams struct { @@ -458,15 +1253,18 @@ type ListResourceTemplatesParams struct { } func (x *ListResourceTemplatesParams) isParams() {} +func (x *ListResourceTemplatesParams) isNil() bool { return x == nil } func (x *ListResourceTemplatesParams) GetProgressToken() any { return getProgressToken(x) } func (x *ListResourceTemplatesParams) SetProgressToken(t any) { setProgressToken(x, t) } func (x *ListResourceTemplatesParams) cursorPtr() *string { return &x.Cursor } // The server's response to a resources/templates/list request from the client. type ListResourceTemplatesResult struct { + completeResultWithType // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. Meta `json:"_meta,omitempty"` + Cacheable // An opaque token representing the pagination position after the last returned // result. If present, there may be more results available. NextCursor string `json:"nextCursor,omitempty"` @@ -474,6 +1272,7 @@ type ListResourceTemplatesResult struct { } func (x *ListResourceTemplatesResult) isResult() {} +func (x *ListResourceTemplatesResult) isNil() bool { return x == nil } func (x *ListResourceTemplatesResult) nextCursorPtr() *string { return &x.NextCursor } type ListResourcesParams struct { @@ -486,15 +1285,18 @@ type ListResourcesParams struct { } func (x *ListResourcesParams) isParams() {} +func (x *ListResourcesParams) isNil() bool { return x == nil } func (x *ListResourcesParams) GetProgressToken() any { return getProgressToken(x) } func (x *ListResourcesParams) SetProgressToken(t any) { setProgressToken(x, t) } func (x *ListResourcesParams) cursorPtr() *string { return &x.Cursor } // The server's response to a resources/list request from the client. type ListResourcesResult struct { + completeResultWithType // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. Meta `json:"_meta,omitempty"` + Cacheable // An opaque token representing the pagination position after the last returned // result. If present, there may be more results available. NextCursor string `json:"nextCursor,omitempty"` @@ -502,8 +1304,15 @@ type ListResourcesResult struct { } func (x *ListResourcesResult) isResult() {} +func (x *ListResourcesResult) isNil() bool { return x == nil } func (x *ListResourcesResult) nextCursorPtr() *string { return &x.NextCursor } +// ListRootsParams holds parameters for a roots/list request. +// +// Deprecated: the roots feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type ListRootsParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. @@ -511,12 +1320,19 @@ type ListRootsParams struct { } func (x *ListRootsParams) isParams() {} +func (x *ListRootsParams) isInputRequest() {} +func (x *ListRootsParams) isNil() bool { return x == nil } func (x *ListRootsParams) GetProgressToken() any { return getProgressToken(x) } func (x *ListRootsParams) SetProgressToken(t any) { setProgressToken(x, t) } // The client's response to a roots/list request from the server. This result // contains an array of Root objects, each representing a root directory or file // that the server can operate on. +// +// Deprecated: the roots feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type ListRootsResult struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. @@ -524,7 +1340,9 @@ type ListRootsResult struct { Roots []*Root `json:"roots"` } -func (*ListRootsResult) isResult() {} +func (*ListRootsResult) isResult() {} +func (x *ListRootsResult) isNil() bool { return x == nil } +func (*ListRootsResult) isInputResponse() {} type ListToolsParams struct { // This property is reserved by the protocol to allow clients and servers to @@ -536,15 +1354,18 @@ type ListToolsParams struct { } func (x *ListToolsParams) isParams() {} +func (x *ListToolsParams) isNil() bool { return x == nil } func (x *ListToolsParams) GetProgressToken() any { return getProgressToken(x) } func (x *ListToolsParams) SetProgressToken(t any) { setProgressToken(x, t) } func (x *ListToolsParams) cursorPtr() *string { return &x.Cursor } // The server's response to a tools/list request from the client. type ListToolsResult struct { + completeResultWithType // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. Meta `json:"_meta,omitempty"` + Cacheable // An opaque token representing the pagination position after the last returned // result. If present, there may be more results available. NextCursor string `json:"nextCursor,omitempty"` @@ -552,14 +1373,27 @@ type ListToolsResult struct { } func (x *ListToolsResult) isResult() {} +func (x *ListToolsResult) isNil() bool { return x == nil } func (x *ListToolsResult) nextCursorPtr() *string { return &x.NextCursor } // The severity of a log message. // // These map to syslog message severities, as specified in RFC-5424: // https://datatracker.ietf.org/doc/html/rfc5424#section-6.2.1 +// +// Deprecated: the logging feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type LoggingLevel string +// LoggingMessageParams holds the parameters for a notifications/message +// notification. +// +// Deprecated: the logging feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type LoggingMessageParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. @@ -574,6 +1408,7 @@ type LoggingMessageParams struct { } func (x *LoggingMessageParams) isParams() {} +func (x *LoggingMessageParams) isNil() bool { return x == nil } func (x *LoggingMessageParams) GetProgressToken() any { return getProgressToken(x) } func (x *LoggingMessageParams) SetProgressToken(t any) { setProgressToken(x, t) } @@ -581,6 +1416,11 @@ func (x *LoggingMessageParams) SetProgressToken(t any) { setProgressToken(x, t) // // Keys not declared here are currently left unspecified by the spec and are up // to the client to interpret. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type ModelHint struct { // A hint for a model name. // @@ -607,6 +1447,11 @@ type ModelHint struct { // These preferences are always advisory. The client may ignore them. It is also // up to the client to decide how to interpret these preferences and how to // balance them against other considerations. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type ModelPreferences struct { // How much to prioritize cost when selecting a model. A value of 0 means cost // is not important, while a value of 1 means cost is the most important factor. @@ -636,6 +1481,7 @@ type PingParams struct { } func (x *PingParams) isParams() {} +func (x *PingParams) isNil() bool { return x == nil } func (x *PingParams) GetProgressToken() any { return getProgressToken(x) } func (x *PingParams) SetProgressToken(t any) { setProgressToken(x, t) } @@ -656,7 +1502,36 @@ type ProgressNotificationParams struct { Total float64 `json:"total,omitempty"` } -func (*ProgressNotificationParams) isParams() {} +func (x *ProgressNotificationParams) isParams() {} +func (x *ProgressNotificationParams) isNil() bool { return x == nil } + +// IconTheme specifies the theme an icon is designed for. +type IconTheme string + +const ( + // IconThemeLight indicates the icon is designed for a light background. + IconThemeLight IconTheme = "light" + // IconThemeDark indicates the icon is designed for a dark background. + IconThemeDark IconTheme = "dark" +) + +// Icon provides visual identifiers for their resources, tools, prompts, and implementations +// See [/specification/draft/basic/index#icons] for notes on icons +// +// TODO(iamsurajbobade): update specification url from draft. +type Icon struct { + // Source is A URI pointing to the icon resource (required). This can be: + // - An HTTP/HTTPS URL pointing to an image file + // - A data URI with base64-encoded image data + Source string `json:"src"` + // Optional MIME type if the server's type is missing or generic + MIMEType string `json:"mimeType,omitempty"` + // Optional size specification (e.g., ["48x48"], ["any"] for scalable formats like SVG, or ["48x48", "96x96"] for multiple sizes) + Sizes []string `json:"sizes,omitempty"` + // Optional theme specifier. "light" indicates the icon is designed for a light + // background, "dark" indicates the icon is designed for a dark background. + Theme IconTheme `json:"theme,omitempty"` +} // A prompt or prompt template that the server offers. type Prompt struct { @@ -673,6 +1548,8 @@ type Prompt struct { // Intended for UI and end-user contexts — optimized to be human-readable and // easily understood, even by those unfamiliar with domain-specific terminology. Title string `json:"title,omitempty"` + // Icons for the prompt, if any. + Icons []Icon `json:"icons,omitempty"` } // Describes an argument that a prompt can accept. @@ -696,6 +1573,7 @@ type PromptListChangedParams struct { } func (x *PromptListChangedParams) isParams() {} +func (x *PromptListChangedParams) isNil() bool { return x == nil } func (x *PromptListChangedParams) GetProgressToken() any { return getProgressToken(x) } func (x *PromptListChangedParams) SetProgressToken(t any) { setProgressToken(x, t) } @@ -716,7 +1594,7 @@ func (m *PromptMessage) UnmarshalJSON(data []byte) error { msg Content *wireContent `json:"content"` } - if err := json.Unmarshal(data, &wire); err != nil { + if err := internaljson.Unmarshal(data, &wire); err != nil { return err } var err error @@ -734,9 +1612,16 @@ type ReadResourceParams struct { // The URI of the resource to read. The URI can use any protocol; it is up to // the server how to interpret it. URI string `json:"uri"` + + // InputResponses maps input request IDs to responses. Set it when the + // client retries a call that asked for input. See [CallToolResult.NeedsInput]. + InputResponses InputResponseMap `json:"inputResponses,omitempty"` + // RequestState is the opaque state from the previous input-required result. + RequestState string `json:"requestState,omitempty"` } func (x *ReadResourceParams) isParams() {} +func (x *ReadResourceParams) isNil() bool { return x == nil } func (x *ReadResourceParams) GetProgressToken() any { return getProgressToken(x) } func (x *ReadResourceParams) SetProgressToken(t any) { setProgressToken(x, t) } @@ -744,11 +1629,70 @@ func (x *ReadResourceParams) SetProgressToken(t any) { setProgressToken(x, t) } type ReadResourceResult struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. - Meta `json:"_meta,omitempty"` + Meta `json:"_meta,omitempty"` + Cacheable Contents []*ResourceContents `json:"contents"` + + // InputRequests is set when the result needs more client input. + // See [CallToolResult.InputRequests]. + InputRequests InputRequestMap `json:"inputRequests,omitempty"` + // RequestState is the opaque state for multi-round-trip retries. + // See [CallToolResult.RequestState]. + RequestState string `json:"requestState,omitempty"` + + // resultType records whether the call finished or needs more client input. + // See [CallToolResult.NeedsInput]. + resultType resultType } -func (*ReadResourceResult) isResult() {} +func (*ReadResourceResult) isResult() {} +func (x *ReadResourceResult) isNil() bool { return x == nil } + +func (r *ReadResourceResult) setResultType(rt resultType) { r.resultType = rt } +func (r *ReadResourceResult) requestState() string { return r.RequestState } +func (r *ReadResourceResult) inputRequests() map[string]InputRequest { + if r == nil { + return nil + } + return r.InputRequests +} +func (r *ReadResourceResult) hasContent() bool { return len(r.Contents) > 0 } + +// NeedsInput reports whether this result requires further client input. +// See [CallToolResult.NeedsInput] for details. +func (r *ReadResourceResult) NeedsInput() bool { return r.resultType == resultTypeInputRequired } + +func (x *ReadResourceResult) MarshalJSON() ([]byte, error) { + type res ReadResourceResult + type wire struct { + res + ResultType resultType `json:"resultType,omitempty"` + InputRequests json.RawMessage `json:"inputRequests,omitempty"` // shadows res.InputRequests + } + w := wire{res: res(*x), ResultType: x.resultType} + if x.InputRequests != nil { + ir, err := json.Marshal(x.InputRequests) + if err != nil { + return nil, err + } + w.InputRequests = ir + } + return json.Marshal(w) +} + +func (x *ReadResourceResult) UnmarshalJSON(data []byte) error { + type res ReadResourceResult + var wire struct { + res + ResultType resultType `json:"resultType"` + } + if err := internaljson.Unmarshal(data, &wire); err != nil { + return err + } + wire.res.resultType = wire.ResultType + *x = ReadResourceResult(wire.res) + return nil +} // A known resource that the server is capable of reading. type Resource struct { @@ -782,6 +1726,8 @@ type Resource struct { Title string `json:"title,omitempty"` // The URI of this resource. URI string `json:"uri"` + // Icons for the resource, if any. + Icons []Icon `json:"icons,omitempty"` } type ResourceListChangedParams struct { @@ -791,6 +1737,7 @@ type ResourceListChangedParams struct { } func (x *ResourceListChangedParams) isParams() {} +func (x *ResourceListChangedParams) isNil() bool { return x == nil } func (x *ResourceListChangedParams) GetProgressToken() any { return getProgressToken(x) } func (x *ResourceListChangedParams) SetProgressToken(t any) { setProgressToken(x, t) } @@ -822,12 +1769,19 @@ type ResourceTemplate struct { // A URI template (according to RFC 6570) that can be used to construct resource // URIs. URITemplate string `json:"uriTemplate"` + // Icons for the resource template, if any. + Icons []Icon `json:"icons,omitempty"` } // The sender or recipient of messages and data in a conversation. type Role string // Represents a root directory or file that the server can operate on. +// +// Deprecated: the roots feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type Root struct { // See [specification/2025-06-18/basic/index#general-fields] for notes on _meta // usage. @@ -842,6 +1796,13 @@ type Root struct { URI string `json:"uri"` } +// RootsListChangedParams holds parameters for a notifications/roots/list_changed +// notification. +// +// Deprecated: the roots feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type RootsListChangedParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. @@ -849,16 +1810,79 @@ type RootsListChangedParams struct { } func (x *RootsListChangedParams) isParams() {} +func (x *RootsListChangedParams) isNil() bool { return x == nil } func (x *RootsListChangedParams) GetProgressToken() any { return getProgressToken(x) } func (x *RootsListChangedParams) SetProgressToken(t any) { setProgressToken(x, t) } -// SamplingCapabilities describes the capabilities for sampling. -type SamplingCapabilities struct{} +// TODO: to be consistent with ServerCapabilities, move the capability types +// below directly above ClientCapabilities. + +// SamplingCapabilities describes the client's support for sampling. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +type SamplingCapabilities struct { + // Context indicates the client supports includeContext values other than "none". + Context *SamplingContextCapabilities `json:"context,omitempty"` + // Tools indicates the client supports tools and toolChoice in sampling requests. + Tools *SamplingToolsCapabilities `json:"tools,omitempty"` +} + +// SamplingContextCapabilities indicates the client supports context inclusion. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +type SamplingContextCapabilities struct{} + +// SamplingToolsCapabilities indicates the client supports tool use in sampling. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +type SamplingToolsCapabilities struct{} + +// ToolChoice controls how the model uses tools during sampling. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +type ToolChoice struct { + // Mode controls tool invocation behavior: + // - "auto": Model decides whether to use tools (default) + // - "required": Model must use at least one tool + // - "none": Model must not use any tools + Mode string `json:"mode,omitempty"` +} // ElicitationCapabilities describes the capabilities for elicitation. -type ElicitationCapabilities struct{} +// +// If neither Form nor URL is set, the 'Form' capability is assumed. +type ElicitationCapabilities struct { + Form *FormElicitationCapabilities `json:"form,omitempty"` + URL *URLElicitationCapabilities `json:"url,omitempty"` +} + +// FormElicitationCapabilities describes capabilities for form elicitation. +type FormElicitationCapabilities struct{} + +// URLElicitationCapabilities describes capabilities for url elicitation. +type URLElicitationCapabilities struct{} // Describes a message issued to or received from an LLM API. +// +// For assistant messages, Content may be text, image, audio, or tool_use. +// For user messages, Content may be text, image, audio, or tool_result. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type SamplingMessage struct { Content Content `json:"content"` Role Role `json:"role"` @@ -872,17 +1896,24 @@ func (m *SamplingMessage) UnmarshalJSON(data []byte) error { msg Content *wireContent `json:"content"` } - if err := json.Unmarshal(data, &wire); err != nil { + if err := internaljson.Unmarshal(data, &wire); err != nil { return err } + // Allow text, image, audio, tool_use, and tool_result in sampling messages var err error - if wire.msg.Content, err = contentFromWire(wire.Content, map[string]bool{"text": true, "image": true, "audio": true}); err != nil { + if wire.msg.Content, err = contentFromWire(wire.Content, map[string]bool{"text": true, "image": true, "audio": true, "tool_use": true, "tool_result": true}); err != nil { return err } *m = SamplingMessage(wire.msg) return nil } +// SetLoggingLevelParams holds parameters for a logging/setLevel request. +// +// Deprecated: the logging feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type SetLoggingLevelParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. @@ -894,6 +1925,7 @@ type SetLoggingLevelParams struct { } func (x *SetLoggingLevelParams) isParams() {} +func (x *SetLoggingLevelParams) isNil() bool { return x == nil } func (x *SetLoggingLevelParams) GetProgressToken() any { return getProgressToken(x) } func (x *SetLoggingLevelParams) SetProgressToken(t any) { setProgressToken(x, t) } @@ -948,8 +1980,17 @@ type Tool struct { // If not provided, Annotations.Title should be used for display if present, // otherwise Name. Title string `json:"title,omitempty"` + // Icons for the tool, if any. + Icons []Icon `json:"icons,omitempty"` } +// hintomitempty is a compatibility parameter that restores the pre-1.7.0 +// behavior of [ToolAnnotations] JSON marshaling, where false-valued bare bool +// fields (ReadOnlyHint, IdempotentHint) were omitted from the output. +// See the documentation for the mcpgodebug package for instructions on how to +// enable it. +var hintomitempty = mcpgodebug.Value("hintomitempty") + // Additional properties describing a Tool to clients. // // NOTE: all properties in ToolAnnotations are hints. They are not @@ -972,7 +2013,7 @@ type ToolAnnotations struct { // (This property is meaningful only when ReadOnlyHint == false.) // // Default: false - IdempotentHint bool `json:"idempotentHint,omitempty"` + IdempotentHint bool `json:"idempotentHint"` // If true, this tool may interact with an "open world" of external entities. If // false, the tool's domain of interaction is closed. For example, the world of // a web search tool is open, whereas that of a memory tool is not. @@ -982,11 +2023,30 @@ type ToolAnnotations struct { // If true, the tool does not modify its environment. // // Default: false - ReadOnlyHint bool `json:"readOnlyHint,omitempty"` + ReadOnlyHint bool `json:"readOnlyHint"` // A human-readable title for the tool. Title string `json:"title,omitempty"` } +// MarshalJSON implements [json.Marshaler] for ToolAnnotations. +// +// To restore the previous behavior where false-valued ReadOnlyHint and +// IdempotentHint were omitted, set MCPGODEBUG=hintomitempty=1. +func (t ToolAnnotations) MarshalJSON() ([]byte, error) { + if hintomitempty == "1" { + type compat struct { + DestructiveHint *bool `json:"destructiveHint,omitempty"` + IdempotentHint bool `json:"idempotentHint,omitempty"` + OpenWorldHint *bool `json:"openWorldHint,omitempty"` + ReadOnlyHint bool `json:"readOnlyHint,omitempty"` + Title string `json:"title,omitempty"` + } + return json.Marshal(compat(t)) + } + type nomethod ToolAnnotations + return json.Marshal(nomethod(t)) +} + type ToolListChangedParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. @@ -994,6 +2054,7 @@ type ToolListChangedParams struct { } func (x *ToolListChangedParams) isParams() {} +func (x *ToolListChangedParams) isNil() bool { return x == nil } func (x *ToolListChangedParams) GetProgressToken() any { return getProgressToken(x) } func (x *ToolListChangedParams) SetProgressToken(t any) { setProgressToken(x, t) } @@ -1007,7 +2068,8 @@ type SubscribeParams struct { URI string `json:"uri"` } -func (*SubscribeParams) isParams() {} +func (x *SubscribeParams) isParams() {} +func (x *SubscribeParams) isNil() bool { return x == nil } // Sent from the client to request cancellation of resources/updated // notifications from the server. This should follow a previous @@ -1020,7 +2082,8 @@ type UnsubscribeParams struct { URI string `json:"uri"` } -func (*UnsubscribeParams) isParams() {} +func (x *UnsubscribeParams) isParams() {} +func (x *UnsubscribeParams) isNil() bool { return x == nil } // A notification from the server to the client, informing it that a resource // has changed and may need to be read again. This should only be sent if the @@ -1033,7 +2096,64 @@ type ResourceUpdatedNotificationParams struct { URI string `json:"uri"` } -func (*ResourceUpdatedNotificationParams) isParams() {} +func (x *ResourceUpdatedNotificationParams) isParams() {} +func (x *ResourceUpdatedNotificationParams) isNil() bool { return x == nil } + +// NotificationSubscriptions describes the set of server-to-client +// notifications a client wishes to receive on a [SubscriptionsListenParams] +// stream. Each field is an explicit opt-in: a server MUST NOT push +// notifications of a type the client did not request. +type NotificationSubscriptions struct { + // ToolsListChanged opts in to "notifications/tools/list_changed". + ToolsListChanged bool `json:"toolsListChanged,omitempty"` + // PromptsListChanged opts in to "notifications/prompts/list_changed". + PromptsListChanged bool `json:"promptsListChanged,omitempty"` + // ResourcesListChanged opts in to "notifications/resources/list_changed". + ResourcesListChanged bool `json:"resourcesListChanged,omitempty"` + // ResourceSubscriptions enumerates the resource URIs for which the client + // wants "notifications/resources/updated". Replaces the legacy + // resources/subscribe RPC. + ResourceSubscriptions []string `json:"resourceSubscriptions,omitempty"` +} + +// SubscriptionsListenParams are the parameters for the +// "subscriptions/listen" RPC. +type SubscriptionsListenParams struct { + // Meta carries the per-request `_meta` triple. + Meta `json:"_meta,omitempty"` + // Notifications declares which notification types the client wants to + // receive on this stream. + Notifications *NotificationSubscriptions `json:"notifications"` +} + +func (x *SubscriptionsListenParams) isParams() {} +func (x *SubscriptionsListenParams) isNil() bool { return x == nil } + +// SubscriptionsAcknowledgedParams are the parameters for the +// "notifications/subscriptions/acknowledged" notification, which the server +// MUST send as the first message on a subscriptions/listen stream. It carries +// the subset of the requested [NotificationSubscriptions] that the server has +// agreed to honor. +type SubscriptionsAcknowledgedParams struct { + Meta `json:"_meta,omitempty"` + Notifications NotificationSubscriptions `json:"notifications"` +} + +func (x *SubscriptionsAcknowledgedParams) isParams() {} +func (x *SubscriptionsAcknowledgedParams) isNil() bool { return x == nil } + +// SubscriptionsListenResult is the response to a "subscriptions/listen" +// request, signalling that the subscription has ended gracefully (for example, +// during server shutdown). Because the listen stream is long-lived, this +// result is sent only when the server tears the subscription down; an abrupt +// transport close carries no response. +type SubscriptionsListenResult struct { + completeResultWithType + Meta `json:"_meta"` +} + +func (*SubscriptionsListenResult) isResult() {} +func (x *SubscriptionsListenResult) isNil() bool { return x == nil } // TODO(jba): add CompleteRequest and related types. @@ -1042,6 +2162,10 @@ type ElicitParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. Meta `json:"_meta,omitempty"` + // The mode of elicitation to use. + // + // If unset, will be inferred from the other fields. + Mode string `json:"mode"` // The message to present to the user. Message string `json:"message"` // A JSON schema object defining the requested elicitation schema. @@ -1055,14 +2179,41 @@ type ElicitParams struct { // map[string]any). // // Only top-level properties are allowed, without nesting. - RequestedSchema any `json:"requestedSchema"` + // + // This is only used for "form" elicitation. + RequestedSchema any `json:"requestedSchema,omitempty"` + // The URL to present to the user. + // + // This is only used for "url" elicitation. + URL string `json:"url,omitempty"` + // The ID of the elicitation. + // + // This is only used for "url" elicitation. + ElicitationID string `json:"elicitationId,omitempty"` } -func (x *ElicitParams) isParams() {} +func (x *ElicitParams) isParams() {} +func (x *ElicitParams) isInputRequest() {} +func (x *ElicitParams) isNil() bool { return x == nil } func (x *ElicitParams) GetProgressToken() any { return getProgressToken(x) } func (x *ElicitParams) SetProgressToken(t any) { setProgressToken(x, t) } +// inferElicitMode returns x with Mode populated by inference if it was empty. +// Mode is inferred as "url" when URL or ElicitationID is set, otherwise "form". +func (x *ElicitParams) inferElicitMode() *ElicitParams { + if x == nil || x.Mode != "" { + return x + } + x2 := *x + if x.URL != "" || x.ElicitationID != "" { + x2.Mode = "url" + } else { + x2.Mode = "form" + } + return &x2 +} + // The client's response to an elicitation/create request from the server. type ElicitResult struct { // This property is reserved by the protocol to allow clients and servers to @@ -1078,70 +2229,143 @@ type ElicitResult struct { Content map[string]any `json:"content,omitempty"` } -func (*ElicitResult) isResult() {} +func (*ElicitResult) isResult() {} +func (x *ElicitResult) isNil() bool { return x == nil } +func (*ElicitResult) isInputResponse() {} -// An Implementation describes the name and version of an MCP implementation, with an optional -// title for UI representation. +// ElicitationCompleteParams is sent from the server to the client, informing it that an out-of-band elicitation interaction has completed. +type ElicitationCompleteParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta `json:"_meta,omitempty"` + // The ID of the elicitation that has completed. This must correspond to the + // elicitationId from the original elicitation/create request. + ElicitationID string `json:"elicitationId"` +} + +func (x *ElicitationCompleteParams) isParams() {} +func (x *ElicitationCompleteParams) isNil() bool { return x == nil } + +// An Implementation describes the name and version of an MCP implementation, with +// optional display metadata. type Implementation struct { // Intended for programmatic or logical use, but used as a display name in past // specs or fallback (if title isn't present). Name string `json:"name"` // Intended for UI and end-user contexts — optimized to be human-readable and // easily understood, even by those unfamiliar with domain-specific terminology. - Title string `json:"title,omitempty"` - Version string `json:"version"` + Title string `json:"title,omitempty"` + // A human-readable description of the implementation. + Description string `json:"description,omitempty"` + Version string `json:"version"` + // WebsiteURL for the server, if any. + WebsiteURL string `json:"websiteUrl,omitempty"` + // Icons for the Server, if any. + Icons []Icon `json:"icons,omitempty"` } -// Present if the server supports argument autocompletion suggestions. +// CompletionCapabilities describes the server's support for argument autocompletion. type CompletionCapabilities struct{} -// Present if the server supports sending log messages to the client. +// LoggingCapabilities describes the server's support for sending log messages to the client. +// +// Deprecated: the logging feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. type LoggingCapabilities struct{} -// Present if the server offers any prompt templates. +// PromptCapabilities describes the server's support for prompts. type PromptCapabilities struct { // Whether this server supports notifications for changes to the prompt list. ListChanged bool `json:"listChanged,omitempty"` } -// Present if the server offers any resources to read. +// ResourceCapabilities describes the server's support for resources. type ResourceCapabilities struct { - // Whether this server supports notifications for changes to the resource list. + // ListChanged reports whether this server supports notifications for + // changes to the resource list. ListChanged bool `json:"listChanged,omitempty"` - // Whether this server supports subscribing to resource updates. + // Subscribe reports whether this server supports subscribing to resource + // updates. Subscribe bool `json:"subscribe,omitempty"` } -// Capabilities that a server may support. Known capabilities are defined here, -// in this schema, but this is not a closed set: any server can define its own, -// additional capabilities. +// ToolCapabilities describes the server's support for tools. +type ToolCapabilities struct { + // ListChanged reports whether this server supports notifications for + // changes to the tool list. + ListChanged bool `json:"listChanged,omitempty"` +} + +// ServerCapabilities describes capabilities that a server supports. type ServerCapabilities struct { - // Present if the server supports argument autocompletion suggestions. - Completions *CompletionCapabilities `json:"completions,omitempty"` - // Experimental, non-standard capabilities that the server supports. + // NOTE: any addition to ServerCapabilities must also be reflected in + // [ServerCapabilities.clone]. + + // Experimental reports non-standard capabilities that the server supports. + // The caller should not modify the map after assigning it. Experimental map[string]any `json:"experimental,omitempty"` - // Present if the server supports sending log messages to the client. + // Extensions reports extensions that the server supports. + // Keys are extension identifiers in "{vendor-prefix}/{extension-name}" format. + // Values are per-extension settings objects; use [ServerCapabilities.AddExtension] + // to ensure nil settings are normalized to empty objects. + // The caller should not modify the map or its values after assigning it. + Extensions map[string]any `json:"extensions,omitempty"` + // Completions is present if the server supports argument autocompletion + // suggestions. + Completions *CompletionCapabilities `json:"completions,omitempty"` + // Logging is present if the server supports log messages. + // + // Deprecated: the logging feature is deprecated as of protocol version + // 2026-07-28 (SEP-2577). It remains functional during the deprecation + // window (at least twelve months). See + // https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. Logging *LoggingCapabilities `json:"logging,omitempty"` - // Present if the server offers any prompt templates. + // Prompts is present if the server supports prompts. Prompts *PromptCapabilities `json:"prompts,omitempty"` - // Present if the server offers any resources to read. + // Resources is present if the server supports resources. Resources *ResourceCapabilities `json:"resources,omitempty"` - // Present if the server offers any tools to call. + // Tools is present if the supports tools. Tools *ToolCapabilities `json:"tools,omitempty"` } -// Present if the server offers any tools to call. -type ToolCapabilities struct { - // Whether this server supports notifications for changes to the tool list. - ListChanged bool `json:"listChanged,omitempty"` +// AddExtension adds an extension with the given name and settings. +// If settings is nil, an empty map is used to ensure valid JSON serialization +// (the spec requires an object, not null). +// The settings map should not be modified after the call. +func (c *ServerCapabilities) AddExtension(name string, settings map[string]any) { + if c.Extensions == nil { + c.Extensions = make(map[string]any) + } + if settings == nil { + settings = map[string]any{} + } + c.Extensions[name] = settings +} + +// clone returns a copy of the ServerCapabilities. +// Values in the Extensions and Experimental maps are shallow-copied. +func (c *ServerCapabilities) clone() *ServerCapabilities { + cp := *c + cp.Experimental = maps.Clone(c.Experimental) + cp.Extensions = maps.Clone(c.Extensions) + cp.Completions = shallowClone(c.Completions) + cp.Logging = shallowClone(c.Logging) + cp.Prompts = shallowClone(c.Prompts) + cp.Resources = shallowClone(c.Resources) + cp.Tools = shallowClone(c.Tools) + return &cp } const ( methodCallTool = "tools/call" notificationCancelled = "notifications/cancelled" methodComplete = "completion/complete" + methodDiscover = "server/discover" methodCreateMessage = "sampling/createMessage" methodElicit = "elicitation/create" + notificationElicitationComplete = "notifications/elicitation/complete" methodGetPrompt = "prompts/get" methodInitialize = "initialize" notificationInitialized = "notifications/initialized" @@ -1160,6 +2384,60 @@ const ( notificationRootsListChanged = "notifications/roots/list_changed" methodSetLevel = "logging/setLevel" methodSubscribe = "resources/subscribe" + methodSubscriptionsListen = "subscriptions/listen" notificationToolListChanged = "notifications/tools/list_changed" methodUnsubscribe = "resources/unsubscribe" + notificationSubscriptionsAck = "notifications/subscriptions/acknowledged" ) + +// Per-request _meta field names for the >= 2026-07-28 protocol version. +// +// These keys appear inside a Params._meta map and carry information that +// previously came from the initialization handshake (SEP-2575). +const ( + // MetaKeyProtocolVersion identifies the MCP protocol version that the + // request follows. + MetaKeyProtocolVersion = "io.modelcontextprotocol/protocolVersion" + // MetaKeyClientInfo carries the client's [Implementation]. + MetaKeyClientInfo = "io.modelcontextprotocol/clientInfo" + // MetaKeyServerInfo carries the server's [Implementation] on responses. + MetaKeyServerInfo = "io.modelcontextprotocol/serverInfo" + // MetaKeyClientCapabilities carries the client's [ClientCapabilities]. + MetaKeyClientCapabilities = "io.modelcontextprotocol/clientCapabilities" + // MetaKeyLogLevel identifies the desired log level for the request. + // + // Deprecated: the logging feature is deprecated as of protocol version + // 2026-07-28 (SEP-2577). It remains functional during the deprecation + // window (at least twelve months). See + // https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. + MetaKeyLogLevel = "io.modelcontextprotocol/logLevel" + // MetaKeySubscriptionID identifies the subscriptions/listen request that an + // out-of-band notification belongs to. + MetaKeySubscriptionID = "io.modelcontextprotocol/subscriptionId" +) + +// UnsupportedProtocolVersionData is the SEP-2575 payload carried in the +// `data` field of a JSON-RPC error response with code +// [CodeUnsupportedProtocolVersion]. The server uses it to advertise which +// versions it supports so the client can pick a mutually supported one. +type UnsupportedProtocolVersionData struct { + // Supported is the list of protocol versions the server supports. + Supported []string `json:"supported"` + // Requested is the protocol version the client asked for. + Requested string `json:"requested"` +} + +// MissingRequiredClientCapabilityData is the SEP-2575 payload carried in the +// `data` field of a JSON-RPC error response with code +// [CodeMissingRequiredClientCapabilities]. The server uses it to indicate +// which client capabilities are required to process the request but were not +// declared by the client in its per-request `_meta` field. +// +// Handlers that require a specific client capability should inspect the +// per-request [ServerRequest.ClientCapabilities] and return a JSON-RPC error +// populated with this structure when the required capability is missing. +type MissingRequiredClientCapabilityData struct { + // RequiredCapabilities is the set of capabilities the server requires + // from the client to process the request. + RequiredCapabilities *ClientCapabilities `json:"requiredCapabilities"` +} diff --git a/mcp/protocol_test.go b/mcp/protocol_test.go index 67d021d..4f0c345 100644 --- a/mcp/protocol_test.go +++ b/mcp/protocol_test.go @@ -10,6 +10,7 @@ import ( "testing" "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" ) func TestParamsMeta(t *testing.T) { @@ -64,6 +65,16 @@ func TestParamsMeta(t *testing.T) { t.Errorf("got %v, want `t`", g) } + // The GetProgressToken and SetProgressToken methods work on a params struct that doesn't have a Meta field. + if g := p2.GetProgressToken(); g != nil { + t.Errorf("got %v, want nil", g) + } + + p2.SetProgressToken("t") + if g := p2.GetProgressToken(); g != "t" { + t.Errorf("got %v, want `t`", g) + } + // You can set a progress token to an int, int32 or int64. p.SetProgressToken(int(1)) p.SetProgressToken(int32(1)) @@ -481,13 +492,768 @@ func TestCompleteResult(t *testing.T) { if err := json.Unmarshal([]byte(test.in), &got); err != nil { t.Fatalf("json.Unmarshal(CompleteResult) failed: %v", err) } - if diff := cmp.Diff(test.want, got); diff != "" { + if diff := cmp.Diff(test.want, got, cmpopts.IgnoreUnexported(CompleteResult{})); diff != "" { t.Errorf("CompleteResult unmarshal mismatch (-want +got):\n%s", diff) } }) } } +// TODO: merge the following 4 tests into content_test.go. +func TestToolUseContent_MarshalJSON(t *testing.T) { + tests := []struct { + name string + content *ToolUseContent + want string + }{ + { + name: "basic tool use", + content: &ToolUseContent{ + ID: "tool_123", + Name: "calculator", + Input: map[string]any{ + "operation": "add", + "x": 1.0, + "y": 2.0, + }, + }, + want: `{"type":"tool_use","id":"tool_123","name":"calculator","input":{"operation":"add","x":1,"y":2}}`, + }, + { + name: "nil input marshals as empty object", + content: &ToolUseContent{ + ID: "tool_456", + Name: "no_args_tool", + Input: nil, + }, + want: `{"type":"tool_use","id":"tool_456","name":"no_args_tool","input":{}}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := tt.content.MarshalJSON() + if err != nil { + t.Fatalf("MarshalJSON() error = %v", err) + } + if diff := cmp.Diff(tt.want, string(got)); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestToolUseContent_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + json string + want *ToolUseContent + }{ + { + name: "basic tool use", + json: `{"type":"tool_use","id":"tool_123","name":"calculator","input":{"x":1,"y":2}}`, + want: &ToolUseContent{ + ID: "tool_123", + Name: "calculator", + Input: map[string]any{ + "x": 1.0, + "y": 2.0, + }, + }, + }, + { + name: "with meta", + json: `{"type":"tool_use","id":"t1","name":"calc","input":{"x":1},"_meta":{"requestId":"req-123"}}`, + want: &ToolUseContent{ + ID: "t1", + Name: "calc", + Input: map[string]any{"x": 1.0}, + Meta: Meta{"requestId": "req-123"}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + wire := &wireContent{} + if err := json.Unmarshal([]byte(tt.json), wire); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + got, err := contentFromWire(wire, map[string]bool{"tool_use": true}) + if err != nil { + t.Fatalf("contentFromWire() error = %v", err) + } + if diff := cmp.Diff(tt.want, got); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestToolResultContent_MarshalJSON(t *testing.T) { + tests := []struct { + name string + content *ToolResultContent + want string + }{ + { + name: "basic tool result", + content: &ToolResultContent{ + ToolUseID: "tool_123", + Content: []Content{&TextContent{Text: "42"}}, + }, + want: `{"type":"tool_result","toolUseId":"tool_123","content":[{"type":"text","text":"42"}]}`, + }, + { + name: "tool result with error", + content: &ToolResultContent{ + ToolUseID: "tool_456", + Content: []Content{&TextContent{Text: "division by zero"}}, + IsError: true, + }, + want: `{"type":"tool_result","toolUseId":"tool_456","content":[{"type":"text","text":"division by zero"}],"isError":true}`, + }, + { + name: "tool result with structured content", + content: &ToolResultContent{ + ToolUseID: "tool_789", + Content: []Content{&TextContent{Text: `{"result": 42}`}}, + StructuredContent: map[string]any{"result": 42.0}, + }, + want: `{"type":"tool_result","toolUseId":"tool_789","content":[{"type":"text","text":"{\"result\": 42}"}],"structuredContent":{"result":42}}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := tt.content.MarshalJSON() + if err != nil { + t.Fatalf("MarshalJSON() error = %v", err) + } + if diff := cmp.Diff(tt.want, string(got)); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestToolResultContent_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + json string + want *ToolResultContent + }{ + { + name: "basic tool result", + json: `{"type":"tool_result","toolUseId":"tool_123","content":[{"type":"text","text":"42"}],"isError":false}`, + want: &ToolResultContent{ + ToolUseID: "tool_123", + Content: []Content{&TextContent{Text: "42"}}, + IsError: false, + }, + }, + { + name: "image nested content", + json: `{"type":"tool_result","toolUseId":"t1","content":[{"type":"image","mimeType":"image/png","data":"YWJj"}]}`, + want: &ToolResultContent{ + ToolUseID: "t1", + Content: []Content{&ImageContent{MIMEType: "image/png", Data: []byte("abc")}}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + wire := &wireContent{} + if err := json.Unmarshal([]byte(tt.json), wire); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + got, err := contentFromWire(wire, map[string]bool{"tool_result": true}) + if err != nil { + t.Fatalf("contentFromWire() error = %v", err) + } + if diff := cmp.Diff(tt.want, got); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestSamplingMessage_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + json string + want *SamplingMessage + }{ + { + name: "tool_use content", + json: `{"content":{"type":"tool_use","id":"tool_1","name":"calc","input":{}},"role":"assistant"}`, + want: &SamplingMessage{ + Role: "assistant", + Content: &ToolUseContent{ + ID: "tool_1", + Name: "calc", + Input: map[string]any{}, + }, + }, + }, + { + name: "tool_result content", + json: `{"content":{"type":"tool_result","toolUseId":"tool_1","content":[{"type":"text","text":"42"}]},"role":"user"}`, + want: &SamplingMessage{ + Role: "user", + Content: &ToolResultContent{ + ToolUseID: "tool_1", + Content: []Content{&TextContent{Text: "42"}}, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got SamplingMessage + if err := json.Unmarshal([]byte(tt.json), &got); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + if diff := cmp.Diff(tt.want, &got); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestSamplingCapabilities_MarshalUnmarshalJSON(t *testing.T) { + tests := []struct { + name string + caps *SamplingCapabilities + json string + }{ + { + name: "WithCapabilities", + caps: &SamplingCapabilities{ + Tools: &SamplingToolsCapabilities{}, + Context: &SamplingContextCapabilities{}, + }, + json: `{"context":{},"tools":{}}`, + }, + { + name: "EmptyCapabilities", + caps: &SamplingCapabilities{}, + json: `{}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotJson, err := json.Marshal(tt.caps) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + if diff := cmp.Diff(tt.json, string(gotJson)); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + var gotCaps SamplingCapabilities + if err := json.Unmarshal([]byte(tt.json), &gotCaps); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + if diff := cmp.Diff(tt.caps, &gotCaps); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestCreateMessageWithToolsParams_MarshalUnmarshalJSON(t *testing.T) { + params := &CreateMessageWithToolsParams{ + MaxTokens: 1000, + Messages: []*SamplingMessageV2{ + { + Role: "user", + Content: []Content{&TextContent{Text: "Calculate 1+1"}}, + }, + }, + Tools: []*Tool{ + { + Name: "calculator", + Description: "A calculator tool", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "x": map[string]any{"type": "number"}, + "y": map[string]any{"type": "number"}, + }, + }, + }, + }, + ToolChoice: &ToolChoice{Mode: "auto"}, + } + + data, err := json.Marshal(params) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + + var got CreateMessageWithToolsParams + if err := json.Unmarshal(data, &got); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + + if diff := cmp.Diff(params, &got); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } +} + +func TestToolChoice_MarshalUnmarshalJSON(t *testing.T) { + choices := []ToolChoice{ + {Mode: "auto"}, + {Mode: "required"}, + {Mode: "none"}, + } + + for _, tc := range choices { + data, err := json.Marshal(tc) + if err != nil { + t.Fatalf("Marshal(%v) error = %v", tc.Mode, err) + } + var got ToolChoice + if err := json.Unmarshal(data, &got); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + if diff := cmp.Diff(tc, got); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + } +} + +func TestCreateMessageWithToolsResult_MarshalJSON(t *testing.T) { + tests := []struct { + name string + result *CreateMessageWithToolsResult + want string + }{ + { + name: "single element content", + result: &CreateMessageWithToolsResult{ + Model: "test", + Role: "assistant", + Content: []Content{&TextContent{Text: "hello"}}, + }, + // Single-element Content marshals as object (not array) for backward compat. + want: `{"content":{"type":"text","text":"hello"},"model":"test","role":"assistant"}`, + }, + { + name: "multiple elements content", + result: &CreateMessageWithToolsResult{ + Model: "test", + Role: "assistant", + Content: []Content{ + &TextContent{Text: "thinking..."}, + &ToolUseContent{ + ID: "call_1", + Name: "calculator", + Input: map[string]any{ + "a": 1.0, + "b": 2.0, + }, + }, + }, + }, + // Multiple elements marshal as array. + want: `{"content":[{"type":"text","text":"thinking..."},{"type":"tool_use","id":"call_1","name":"calculator","input":{"a":1,"b":2}}],"model":"test","role":"assistant"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := json.Marshal(tt.result) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + if diff := cmp.Diff(tt.want, string(got)); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestCreateMessageWithToolsResult_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + json string + want *CreateMessageWithToolsResult + wantErr bool + }{ + { + name: "single tool_use content", + json: `{"content":{"type":"tool_use","id":"tool_1","name":"calculator","input":{"x":1}},"model":"test-model","role":"assistant","stopReason":"toolUse"}`, + want: &CreateMessageWithToolsResult{ + Model: "test-model", + Role: "assistant", + StopReason: "toolUse", + Content: []Content{ + &ToolUseContent{ + ID: "tool_1", + Name: "calculator", + Input: map[string]any{ + "x": 1.0, + }, + }, + }, + }, + }, + { + name: "array of tool_use content", + json: `{"content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}},{"type":"tool_use","id":"t2","name":"search","input":{"q":"hi"}}],"model":"test","role":"assistant","stopReason":"toolUse"}`, + want: &CreateMessageWithToolsResult{ + Model: "test", + Role: "assistant", + StopReason: "toolUse", + Content: []Content{ + &ToolUseContent{ + ID: "t1", + Name: "calc", + Input: map[string]any{"x": 1.0}, + }, + &ToolUseContent{ + ID: "t2", + Name: "search", + Input: map[string]any{"q": "hi"}, + }, + }, + }, + }, + { + name: "empty array", + json: `{"content":[],"model":"m","role":"assistant"}`, + want: &CreateMessageWithToolsResult{ + Model: "m", + Role: "assistant", + Content: []Content{}, + }, + }, + { + name: "null content", + json: `{"content":null,"model":"m","role":"assistant"}`, + wantErr: true, + }, + { + name: "rejects tool_result", + json: `{"content":{"type":"tool_result","toolUseId":"t1","content":[]},"model":"m","role":"assistant"}`, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got CreateMessageWithToolsResult + err := json.Unmarshal([]byte(tt.json), &got) + if (err != nil) != tt.wantErr { + t.Fatalf("Unmarshal() error = %v, wantErr %v", err, tt.wantErr) + } + if tt.wantErr { + return + } + if diff := cmp.Diff(tt.want, &got); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestSamplingMessageV2_MarshalJSON(t *testing.T) { + tests := []struct { + name string + msg *SamplingMessageV2 + want string + }{ + { + name: "empty content", + msg: &SamplingMessageV2{ + Role: "user", + Content: []Content{}, + }, + want: `{"content":[],"role":"user"}`, + }, + { + name: "single content", + msg: &SamplingMessageV2{ + Role: "assistant", + Content: []Content{&TextContent{Text: "hello"}}, + }, + want: `{"content":{"type":"text","text":"hello"},"role":"assistant"}`, + }, + { + name: "multiple content", + msg: &SamplingMessageV2{ + Role: "assistant", + // Text + tool_use in the same message (valid per spec for assistant). + Content: []Content{ + &TextContent{Text: "checking weather"}, + &ToolUseContent{ID: "c1", Name: "weather", Input: map[string]any{"city": "SF"}}, + }, + }, + want: `{"content":[{"type":"text","text":"checking weather"},{"type":"tool_use","id":"c1","name":"weather","input":{"city":"SF"}}],"role":"assistant"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := json.Marshal(tt.msg) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + if diff := cmp.Diff(tt.want, string(got)); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestSamplingMessageV2_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + json string + want *SamplingMessageV2 + }{ + { + name: "single content object", + json: `{"role":"user","content":{"type":"text","text":"hello"}}`, + want: &SamplingMessageV2{ + Role: "user", + Content: []Content{ + &TextContent{Text: "hello"}, + }, + }, + }, + { + name: "multiple content", + json: `{"role":"assistant","content":[{"type":"text","text":"Let me check the weather."},{"type":"tool_use","id":"c1","name":"weather","input":{"city":"SF"}}]}`, + want: &SamplingMessageV2{ + Role: "assistant", + // Text + tool_use in the same message (valid per spec for assistant). + Content: []Content{ + &TextContent{Text: "Let me check the weather."}, + &ToolUseContent{ID: "c1", Name: "weather", Input: map[string]any{"city": "SF"}}, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got SamplingMessageV2 + if err := json.Unmarshal([]byte(tt.json), &got); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + if diff := cmp.Diff(tt.want, &got); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestToBase_Conversion(t *testing.T) { + tests := []struct { + name string + params *CreateMessageWithToolsParams + want *CreateMessageParams + wantErr bool + }{ + { + name: "Single content", + params: &CreateMessageWithToolsParams{ + MaxTokens: 1000, + Messages: []*SamplingMessageV2{ + {Role: "user", Content: []Content{&TextContent{Text: "hello"}}}, + }, + Tools: []*Tool{{Name: "calc"}}, + ToolChoice: &ToolChoice{Mode: "auto"}, + }, + want: &CreateMessageParams{ + MaxTokens: 1000, + Messages: []*SamplingMessage{ + {Role: "user", Content: &TextContent{Text: "hello"}}, + }, + }, + }, + { + name: "Multi content", + params: &CreateMessageWithToolsParams{ + MaxTokens: 1000, + Messages: []*SamplingMessageV2{ + {Role: "assistant", Content: []Content{ + &ToolUseContent{ID: "c1", Name: "calc", Input: map[string]any{}}, + &ToolUseContent{ID: "c2", Name: "search", Input: map[string]any{}}, + }}, + }, + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := tt.params.toBase() + if (err != nil) != tt.wantErr { + t.Fatalf("toBase() error = %v, wantErr %v", err, tt.wantErr) + } + if tt.wantErr { + return + } + if diff := cmp.Diff(tt.want, got); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestToWithTools_Conversion(t *testing.T) { + tests := []struct { + name string + result *CreateMessageResult + want *CreateMessageWithToolsResult + }{ + { + name: "with content", + result: &CreateMessageResult{ + Model: "test", + Role: "assistant", + Content: &TextContent{Text: "hello"}, + StopReason: "endTurn", + }, + want: &CreateMessageWithToolsResult{ + Model: "test", + Role: "assistant", + Content: []Content{&TextContent{Text: "hello"}}, + StopReason: "endTurn", + }, + }, + { + name: "nil content", + result: &CreateMessageResult{ + Model: "test", + Role: "assistant", + }, + want: &CreateMessageWithToolsResult{ + Model: "test", + Role: "assistant", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + wt := tt.result.toWithTools() + if diff := cmp.Diff(tt.want, wt); diff != "" { + t.Errorf("mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestInputRequestMapJSON(t *testing.T) { + t.Run("nil is omitted from JSON", func(t *testing.T) { + result := CallToolResult{Content: []Content{&TextContent{Text: "ok"}}} + data, err := json.Marshal(&result) + if err != nil { + t.Fatal(err) + } + var raw map[string]json.RawMessage + if err := json.Unmarshal(data, &raw); err != nil { + t.Fatal(err) + } + if _, ok := raw["inputRequests"]; ok { + t.Errorf("nil InputRequests should be omitted, got %s", raw["inputRequests"]) + } + }) + + t.Run("non-nil empty round-trips", func(t *testing.T) { + result := CallToolResult{ + Content: []Content{&TextContent{Text: "ok"}}, + InputRequests: InputRequestMap{}, + } + data, err := json.Marshal(&result) + if err != nil { + t.Fatal(err) + } + var raw map[string]json.RawMessage + if err := json.Unmarshal(data, &raw); err != nil { + t.Fatal(err) + } + if string(raw["inputRequests"]) != "{}" { + t.Errorf("empty InputRequests should marshal to {}, got %s", raw["inputRequests"]) + } + var got CallToolResult + if err := json.Unmarshal(data, &got); err != nil { + t.Fatal(err) + } + if got.InputRequests == nil { + t.Error("empty InputRequests should round-trip as non-nil") + } + }) + + t.Run("populated round-trips", func(t *testing.T) { + result := CallToolResult{ + Content: []Content{&TextContent{Text: "ok"}}, + InputRequests: InputRequestMap{ + "r1": &ElicitParams{Message: "confirm?"}, + }, + } + data, err := json.Marshal(&result) + if err != nil { + t.Fatal(err) + } + var got CallToolResult + if err := json.Unmarshal(data, &got); err != nil { + t.Fatal(err) + } + if got.InputRequests == nil { + t.Fatal("InputRequests should not be nil after round-trip") + } + if _, ok := got.InputRequests["r1"]; !ok { + t.Error("expected key r1 in InputRequests") + } + }) +} + +func TestInputResponseMapJSON(t *testing.T) { + tests := []struct { + name string + value InputResponse + check func(t *testing.T, got InputResponse) + }{ + { + name: "elicit", + value: &ElicitResult{Action: "accept", Content: map[string]any{"ok": true}}, + }, + { + name: "sampling", + value: &CreateMessageWithToolsResult{Role: "assistant", Model: "test-model", Content: []Content{&TextContent{Text: "hello"}}}, + }, + { + name: "list-roots", + value: &ListRootsResult{Roots: []*Root{{URI: "file:///tmp", Name: "tmp"}}}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + key := tt.name + data, err := json.Marshal(InputResponseMap{key: tt.value}) + if err != nil { + t.Fatal(err) + } + var got InputResponseMap + if err := json.Unmarshal(data, &got); err != nil { + t.Fatal(err) + } + if diff := cmp.Diff(tt.value, got[key], ctrCmpOpts...); diff != "" { + t.Errorf("mismatch (-want, +got):\n%s", diff) + } + }) + } +} + func TestContentUnmarshal(t *testing.T) { // Verify that types with a Content field round-trip properly. roundtrip := func(in, out any) { @@ -519,7 +1285,7 @@ func TestContentUnmarshal(t *testing.T) { Meta: Meta{"m": true}, Content: content, IsError: true, - StructuredContent: 3.0, + StructuredContent: json.Number("3"), } var gotf CallToolResult roundtrip(ctrf, &gotf) @@ -531,3 +1297,78 @@ func TestContentUnmarshal(t *testing.T) { var gotpm PromptMessage roundtrip(pm, &gotpm) } + +func TestCallToolResultStructuredContentFloat64Compatibility(t *testing.T) { + prev := structuredcontentfloat64 + structuredcontentfloat64 = "1" + t.Cleanup(func() { structuredcontentfloat64 = prev }) + + var got CallToolResult + if err := json.Unmarshal([]byte(`{"content":[],"structuredContent":{"id":9007199254740993}}`), &got); err != nil { + t.Fatal(err) + } + structured, ok := got.StructuredContent.(map[string]any) + if !ok { + t.Fatalf("StructuredContent type = %T, want map[string]any", got.StructuredContent) + } + if got := structured["id"]; got != float64(9007199254740992) { + t.Errorf("id = %v (%T), want rounded float64 value", got, got) + } +} + +func TestToolAnnotations_MarshalJSON(t *testing.T) { + boolPtr := func(b bool) *bool { return &b } + + tests := []struct { + name string + in ToolAnnotations + want string + }{ + { + name: "ZeroValue", + in: ToolAnnotations{}, + want: `{"idempotentHint":false,"readOnlyHint":false}`, + }, + { + name: "AllFalse", + in: ToolAnnotations{ + DestructiveHint: boolPtr(false), + IdempotentHint: false, + OpenWorldHint: boolPtr(false), + ReadOnlyHint: false, + }, + want: `{"destructiveHint":false,"idempotentHint":false,"openWorldHint":false,"readOnlyHint":false}`, + }, + { + name: "AllTrue", + in: ToolAnnotations{ + DestructiveHint: boolPtr(true), + IdempotentHint: true, + OpenWorldHint: boolPtr(true), + ReadOnlyHint: true, + Title: "my tool", + }, + want: `{"destructiveHint":true,"idempotentHint":true,"openWorldHint":true,"readOnlyHint":true,"title":"my tool"}`, + }, + { + name: "MixedValues", + in: ToolAnnotations{ + ReadOnlyHint: true, + Title: "read tool", + }, + want: `{"idempotentHint":false,"readOnlyHint":true,"title":"read tool"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := json.Marshal(tt.in) + if err != nil { + t.Fatalf("json.Marshal(%v) failed: %v", tt.in, err) + } + if diff := cmp.Diff(tt.want, string(got)); diff != "" { + t.Errorf("json.Marshal() mismatch (-want +got):\n%s", diff) + } + }) + } +} diff --git a/mcp/requests.go b/mcp/requests.go index 82b700f..64414ca 100644 --- a/mcp/requests.go +++ b/mcp/requests.go @@ -19,19 +19,23 @@ type ( ReadResourceRequest = ServerRequest[*ReadResourceParams] RootsListChangedRequest = ServerRequest[*RootsListChangedParams] SubscribeRequest = ServerRequest[*SubscribeParams] + SubscriptionsListenRequest = ServerRequest[*SubscriptionsListenParams] UnsubscribeRequest = ServerRequest[*UnsubscribeParams] ) type ( - CreateMessageRequest = ClientRequest[*CreateMessageParams] - ElicitRequest = ClientRequest[*ElicitParams] - initializedClientRequest = ClientRequest[*InitializedParams] - InitializeRequest = ClientRequest[*InitializeParams] - ListRootsRequest = ClientRequest[*ListRootsParams] - LoggingMessageRequest = ClientRequest[*LoggingMessageParams] - ProgressNotificationClientRequest = ClientRequest[*ProgressNotificationParams] - PromptListChangedRequest = ClientRequest[*PromptListChangedParams] - ResourceListChangedRequest = ClientRequest[*ResourceListChangedParams] - ResourceUpdatedNotificationRequest = ClientRequest[*ResourceUpdatedNotificationParams] - ToolListChangedRequest = ClientRequest[*ToolListChangedParams] + CreateMessageRequest = ClientRequest[*CreateMessageParams] + CreateMessageWithToolsRequest = ClientRequest[*CreateMessageWithToolsParams] + DiscoverRequest = ClientRequest[*DiscoverParams] + ElicitRequest = ClientRequest[*ElicitParams] + initializedClientRequest = ClientRequest[*InitializedParams] + InitializeRequest = ClientRequest[*InitializeParams] + ListRootsRequest = ClientRequest[*ListRootsParams] + LoggingMessageRequest = ClientRequest[*LoggingMessageParams] + ProgressNotificationClientRequest = ClientRequest[*ProgressNotificationParams] + PromptListChangedRequest = ClientRequest[*PromptListChangedParams] + ResourceListChangedRequest = ClientRequest[*ResourceListChangedParams] + ResourceUpdatedNotificationRequest = ClientRequest[*ResourceUpdatedNotificationParams] + ToolListChangedRequest = ClientRequest[*ToolListChangedParams] + ElicitationCompleteNotificationRequest = ClientRequest[*ElicitationCompleteParams] ) diff --git a/mcp/resource.go b/mcp/resource.go index 8746eda..fff38a2 100644 --- a/mcp/resource.go +++ b/mcp/resource.go @@ -15,8 +15,9 @@ import ( "path/filepath" "strings" - "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/internal/mcpgodebug" "github.com/modelcontextprotocol/go-sdk/internal/util" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" "github.com/yosida95/uritemplate/v3" ) @@ -37,11 +38,28 @@ type serverResourceTemplate struct { // If it cannot find the resource, it should return the result of calling [ResourceNotFoundError]. type ResourceHandler func(context.Context, *ReadResourceRequest) (*ReadResourceResult, error) +// customresnotfounderrcode is a compatibility parameter that restores the +// pre-1.7.0 behavior of [ResourceNotFoundError] and [CodeResourceNotFound], +// where the error code was a custom -32002. See the documentation for the mcpgodebug +// package for instructions on how to enable it. +// The option will be removed in the future version of the SDK. +var customresnotfounderrcode = mcpgodebug.Value("customresnotfounderrcode") + +func init() { + if customresnotfounderrcode == "1" { + CodeResourceNotFound = -32002 + } +} + // ResourceNotFoundError returns an error indicating that a resource being read could // not be found. +// +// By default, the error code is -32602 (Invalid Params), as specified in the +// MCP specification (SEP-2164). To restore the pre-1.7.0 release behavior where the +// error code was -32002, set MCPGODEBUG=customresnotfounderrcode=1. func ResourceNotFoundError(uri string) error { - return &jsonrpc2.WireError{ - Code: codeResourceNotFound, + return &jsonrpc.Error{ + Code: CodeResourceNotFound, Message: "Resource not found", Data: json.RawMessage(fmt.Sprintf(`{"uri":%q}`, uri)), } @@ -113,6 +131,23 @@ func computeURIFilepath(rawURI, dirFilepath string, rootFilepaths []string) (str return uriFilepathRel, nil } +// withFile calls f on the file at join(dir, rel), +// protecting against path traversal attacks. +func withFile(dir, rel string, f func(*os.File) error) (err error) { + r, err := os.OpenRoot(dir) + if err != nil { + return err + } + defer r.Close() + file, err := r.Open(rel) + if err != nil { + return err + } + // Record error, in case f writes. + defer func() { err = errors.Join(err, file.Close()) }() + return f(file) +} + // fileRoots transforms the Roots obtained from the client into absolute paths on // the local filesystem. // TODO(jba): expose this functionality to user ResourceHandlers, diff --git a/mcp/resource_go124.go b/mcp/resource_go124.go deleted file mode 100644 index 4a35603..0000000 --- a/mcp/resource_go124.go +++ /dev/null @@ -1,29 +0,0 @@ -// Copyright 2025 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -//go:build go1.24 - -package mcp - -import ( - "errors" - "os" -) - -// withFile calls f on the file at join(dir, rel), -// protecting against path traversal attacks. -func withFile(dir, rel string, f func(*os.File) error) (err error) { - r, err := os.OpenRoot(dir) - if err != nil { - return err - } - defer r.Close() - file, err := r.Open(rel) - if err != nil { - return err - } - // Record error, in case f writes. - defer func() { err = errors.Join(err, file.Close()) }() - return f(file) -} diff --git a/mcp/resource_pre_go124.go b/mcp/resource_pre_go124.go deleted file mode 100644 index d1f72ee..0000000 --- a/mcp/resource_pre_go124.go +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright 2025 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -//go:build !go1.24 - -package mcp - -import ( - "errors" - "os" - "path/filepath" -) - -// withFile calls f on the file at join(dir, rel). -// It does not protect against path traversal attacks. -func withFile(dir, rel string, f func(*os.File) error) (err error) { - file, err := os.Open(filepath.Join(dir, rel)) - if err != nil { - return err - } - // Record error, in case f writes. - defer func() { err = errors.Join(err, file.Close()) }() - return f(file) -} diff --git a/mcp/sampling_test.go b/mcp/sampling_test.go new file mode 100644 index 0000000..7a3ecab --- /dev/null +++ b/mcp/sampling_test.go @@ -0,0 +1,468 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +//lint:file-ignore SA1019 tests exercise deprecated SEP-2577 sampling APIs. + +package mcp + +// TODO: move other sampling-related tests to this file. + +import ( + "context" + "strings" + "testing" + + "github.com/google/go-cmp/cmp" +) + +func TestSamplingWithTools_ToolUse(t *testing.T) { + ctx := context.Background() + ct, st := NewInMemoryTransports() + + // Track what the client received + var gotParams *CreateMessageWithToolsParams + result := &CreateMessageWithToolsResult{ + Model: "test-model", + Role: "assistant", + Content: []Content{ + &ToolUseContent{ + ID: "tool_call_1", + Name: "calculator", + Input: map[string]any{"x": 1.0, "y": 2.0}, + }, + }, + StopReason: "toolUse", + } + + // Client with tools capability, using CreateMessageWithToolsHandler + client := NewClient(testImpl, &ClientOptions{ + CreateMessageWithToolsHandler: func(_ context.Context, req *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + gotParams = req.Params + return result, nil + }, + Capabilities: &ClientCapabilities{ + Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}}, + }, + }) + + server := NewServer(testImpl, nil) + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Server sends CreateMessageWithTools + params := &CreateMessageWithToolsParams{ + MaxTokens: 1000, + Messages: []*SamplingMessageV2{ + {Role: "user", Content: []Content{&TextContent{Text: "Calculate 1+2"}}}, + }, + Tools: []*Tool{ + { + Name: "calculator", + Description: "A calculator", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "x": map[string]any{"type": "number"}, + "y": map[string]any{"type": "number"}, + }, + }, + }, + }, + ToolChoice: &ToolChoice{Mode: "auto"}, + } + gotResult, err := ss.CreateMessageWithTools(ctx, params) + if err != nil { + t.Fatalf("CreateMessageWithTools() error = %v", err) + } + + // Verify client received the params + if diff := cmp.Diff(params, gotParams); diff != "" { + t.Errorf("CreateMessageWithToolsParams mismatch (-want +got):\n%s", diff) + } + + // Verify server received the tool use response + if diff := cmp.Diff(result, gotResult); diff != "" { + t.Errorf("CreateMessageWithToolsResult mismatch (-want +got):\n%s", diff) + } +} + +func TestSamplingWithTools_ToolResult(t *testing.T) { + ctx := context.Background() + ct, st := NewInMemoryTransports() + + // Track messages received by client + var gotParams *CreateMessageParams + + client := NewClient(testImpl, &ClientOptions{ + CreateMessageHandler: func(_ context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) { + gotParams = req.Params + return &CreateMessageResult{ + Model: "test-model", + Role: "assistant", + Content: &TextContent{Text: "The result is 3"}, + }, nil + }, + Capabilities: &ClientCapabilities{ + Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}}, + }, + }) + + server := NewServer(testImpl, nil) + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + params := &CreateMessageParams{ + MaxTokens: 1000, + Messages: []*SamplingMessage{ + {Role: "user", Content: &TextContent{Text: "Calculate 1+2"}}, + {Role: "assistant", Content: &ToolUseContent{ + ID: "tool_1", + Name: "calculator", + Input: map[string]any{"x": 1.0, "y": 2.0}, + }}, + {Role: "user", Content: &ToolResultContent{ + ToolUseID: "tool_1", + Content: []Content{&TextContent{Text: "3"}}, + }}, + }, + } + _, err = ss.CreateMessage(ctx, params) + if err != nil { + t.Fatalf("CreateMessage() error = %v", err) + } + + if diff := cmp.Diff(params, gotParams); diff != "" { + t.Errorf("CreateMessageParams mismatch (-want +got):\n%s", diff) + } +} + +func TestSamplingToolsCapabilities(t *testing.T) { + ctx := context.Background() + + t.Run("client with explicit tools capability", func(t *testing.T) { + ct, st := NewInMemoryTransports() + + client := NewClient(testImpl, &ClientOptions{ + CreateMessageHandler: func(_ context.Context, _ *CreateMessageRequest) (*CreateMessageResult, error) { + return &CreateMessageResult{Model: "m", Content: &TextContent{}}, nil + }, + Capabilities: &ClientCapabilities{ + Sampling: &SamplingCapabilities{ + Tools: &SamplingToolsCapabilities{}, + Context: &SamplingContextCapabilities{}, + }, + }, + }) + + server := NewServer(testImpl, nil) + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + cs, err := client.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Check server sees client capabilities + caps := ss.InitializeParams().Capabilities + want := &SamplingCapabilities{ + Tools: &SamplingToolsCapabilities{}, + Context: &SamplingContextCapabilities{}, + } + if diff := cmp.Diff(want, caps.Sampling); diff != "" { + t.Errorf("SamplingCapabilities mismatch (-want +got):\n%s", diff) + } + }) + + t.Run("client without tools capability", func(t *testing.T) { + ct, st := NewInMemoryTransports() + + client := NewClient(testImpl, &ClientOptions{ + CreateMessageHandler: func(_ context.Context, _ *CreateMessageRequest) (*CreateMessageResult, error) { + return &CreateMessageResult{Model: "m", Content: &TextContent{}}, nil + }, + // No Capabilities.Sampling.Tools set + }) + + server := NewServer(testImpl, nil) + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + cs, err := client.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Check server sees client capabilities + caps := ss.InitializeParams().Capabilities + want := &SamplingCapabilities{} + if diff := cmp.Diff(want, caps.Sampling); diff != "" { + t.Errorf("SamplingCapabilities mismatch (-want +got):\n%s", diff) + } + }) + + t.Run("client with inferred tools capability", func(t *testing.T) { + ct, st := NewInMemoryTransports() + + client := NewClient(testImpl, &ClientOptions{ + CreateMessageWithToolsHandler: func(_ context.Context, _ *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + return &CreateMessageWithToolsResult{Model: "m", Content: []Content{&TextContent{}}}, nil + }, + // No explicit Capabilities set — tools should be inferred. + }) + + server := NewServer(testImpl, nil) + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + cs, err := client.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + caps := ss.InitializeParams().Capabilities + want := &SamplingCapabilities{ + Tools: &SamplingToolsCapabilities{}, + } + if diff := cmp.Diff(want, caps.Sampling); diff != "" { + t.Errorf("SamplingCapabilities mismatch (-want +got):\n%s", diff) + } + }) +} + +func TestSamplingWithTools_ToolResultWithError(t *testing.T) { + ctx := context.Background() + ct, st := NewInMemoryTransports() + + var gotParams *CreateMessageParams + + client := NewClient(testImpl, &ClientOptions{ + CreateMessageHandler: func(_ context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) { + gotParams = req.Params + return &CreateMessageResult{ + Model: "test-model", + Role: "assistant", + Content: &TextContent{Text: "I see the tool failed"}, + }, nil + }, + Capabilities: &ClientCapabilities{ + Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}}, + }, + }) + + server := NewServer(testImpl, nil) + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + params := &CreateMessageParams{ + MaxTokens: 1000, + Messages: []*SamplingMessage{ + {Role: "user", Content: &TextContent{Text: "Divide 1 by 0"}}, + {Role: "assistant", Content: &ToolUseContent{ + ID: "tool_1", + Name: "calculator", + Input: map[string]any{"op": "div", "x": 1.0, "y": 0.0}, + }}, + {Role: "user", Content: &ToolResultContent{ + ToolUseID: "tool_1", + Content: []Content{&TextContent{Text: "division by zero"}}, + IsError: true, + }}, + }, + } + _, err = ss.CreateMessage(ctx, params) + if err != nil { + t.Fatalf("CreateMessage() error = %v", err) + } + + if diff := cmp.Diff(params, gotParams); diff != "" { + t.Errorf("CreateMessageParams mismatch (-want +got):\n%s", diff) + } +} + +func TestSamplingWithTools_ParallelToolCalls(t *testing.T) { + ctx := context.Background() + ct, st := NewInMemoryTransports() + + result := &CreateMessageWithToolsResult{ + Model: "test-model", + Role: "assistant", + Content: []Content{ + &ToolUseContent{ID: "call_1", Name: "weather", Input: map[string]any{"city": "SF"}}, + &ToolUseContent{ID: "call_2", Name: "weather", Input: map[string]any{"city": "NY"}}, + }, + StopReason: "toolUse", + } + // Client returns parallel tool use results + client := NewClient(testImpl, &ClientOptions{ + CreateMessageWithToolsHandler: func(_ context.Context, req *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + return result, nil + }, + Capabilities: &ClientCapabilities{ + Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}}, + }, + }) + + server := NewServer(testImpl, nil) + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + gotResult, err := ss.CreateMessageWithTools(ctx, &CreateMessageWithToolsParams{ + MaxTokens: 1000, + Messages: []*SamplingMessageV2{ + {Role: "user", Content: []Content{&TextContent{Text: "Weather in SF and NY"}}}, + }, + Tools: []*Tool{ + {Name: "weather", InputSchema: map[string]any{"type": "object"}}, + }, + }) + if err != nil { + t.Fatalf("CreateMessageWithTools() error = %v", err) + } + + if diff := cmp.Diff(result, gotResult); diff != "" { + t.Errorf("CreateMessageWithToolsResult mismatch (-want +got):\n%s", diff) + } +} + +func TestNewClient_BothHandlersPanics(t *testing.T) { + defer func() { + r := recover() + if r == nil { + t.Fatal("expected panic when both handlers set") + } + msg, ok := r.(string) + if !ok || !strings.Contains(msg, "CreateMessageHandler") { + t.Errorf("unexpected panic: %v", r) + } + }() + NewClient(testImpl, &ClientOptions{ + CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + return nil, nil + }, + CreateMessageWithToolsHandler: func(context.Context, *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + return nil, nil + }, + }) +} + +func TestCreateMessage_MultipleContentError(t *testing.T) { + ctx := context.Background() + ct, st := NewInMemoryTransports() + + // Client returns multiple content blocks via CreateMessageWithToolsHandler + client := NewClient(testImpl, &ClientOptions{ + CreateMessageWithToolsHandler: func(_ context.Context, _ *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) { + return &CreateMessageWithToolsResult{ + Model: "test", + Role: "assistant", + Content: []Content{ + &TextContent{Text: "a"}, + &TextContent{Text: "b"}, + }, + }, nil + }, + }) + + server := NewServer(testImpl, nil) + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + cs, err := client.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Server calls CreateMessage (singular), should get error + _, err = ss.CreateMessage(ctx, &CreateMessageParams{ + MaxTokens: 100, + Messages: []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "hi"}}}, + }) + if err == nil { + t.Fatal("expected error for multiple content blocks") + } + if !strings.Contains(err.Error(), "CreateMessageWithTools") { + t.Errorf("error should mention CreateMessageWithTools, got: %v", err) + } +} + +func TestClientCapabilities_CloneSampling(t *testing.T) { + caps := &ClientCapabilities{ + Sampling: &SamplingCapabilities{ + Tools: &SamplingToolsCapabilities{}, + Context: &SamplingContextCapabilities{}, + }, + } + cloned := caps.clone() + + // Verify deep copy — Sampling pointer should differ. + // (Tools and Context are empty structs, so Go may reuse the same address; + // we just check they're non-nil and that mutating Sampling doesn't alias.) + if cloned.Sampling == caps.Sampling { + t.Error("Sampling pointer should differ after clone") + } + if cloned.Sampling.Tools == nil { + t.Error("cloned Sampling.Tools should not be nil") + } + if cloned.Sampling.Context == nil { + t.Error("cloned Sampling.Context should not be nil") + } + // Verify mutation doesn't affect original. + cloned.Sampling.Tools = nil + if caps.Sampling.Tools == nil { + t.Error("modifying cloned Sampling.Tools should not affect original") + } +} diff --git a/mcp/schema_cache.go b/mcp/schema_cache.go new file mode 100644 index 0000000..5fb032d --- /dev/null +++ b/mcp/schema_cache.go @@ -0,0 +1,69 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "reflect" + "sync" + + "github.com/google/jsonschema-go/jsonschema" +) + +// A SchemaCache caches JSON schemas to avoid repeated reflection and resolution. +// +// This is useful for stateless server deployments (one [Server] per request) +// where tools are re-registered on every request. Without caching, each +// [AddTool] call triggers expensive reflection-based schema generation. +// +// A SchemaCache is safe for concurrent use by multiple goroutines. +// +// # Trade-offs +// +// The cache is unbounded: it stores one entry per unique Go type or schema +// pointer. For typical MCP servers with a fixed set of tools, memory usage +// is negligible. However, if tool input types are generated dynamically, +// the cache will grow without bound. +// +// The cache uses pointer identity for pre-defined schemas. If a schema's +// contents change but the pointer remains the same, stale resolved schemas +// may be returned. In practice, this is not an issue because tool schemas +// are typically defined once at startup. +type SchemaCache struct { + byType sync.Map // reflect.Type -> *cachedSchema + bySchema sync.Map // *jsonschema.Schema -> *jsonschema.Resolved +} + +type cachedSchema struct { + schema *jsonschema.Schema + resolved *jsonschema.Resolved +} + +// NewSchemaCache creates a new [SchemaCache]. +func NewSchemaCache() *SchemaCache { + return &SchemaCache{} +} + +func (c *SchemaCache) getByType(t reflect.Type) (*jsonschema.Schema, *jsonschema.Resolved, bool) { + if v, ok := c.byType.Load(t); ok { + cs := v.(*cachedSchema) + return cs.schema, cs.resolved, true + } + return nil, nil, false +} + +func (c *SchemaCache) setByType(t reflect.Type, schema *jsonschema.Schema, resolved *jsonschema.Resolved) { + c.byType.Store(t, &cachedSchema{schema: schema, resolved: resolved}) +} + +func (c *SchemaCache) getBySchema(schema *jsonschema.Schema) (*jsonschema.Resolved, bool) { + if v, ok := c.bySchema.Load(schema); ok { + return v.(*jsonschema.Resolved), true + } + return nil, false +} + +func (c *SchemaCache) setBySchema(schema *jsonschema.Schema, resolved *jsonschema.Resolved) { + c.bySchema.Store(schema, resolved) +} diff --git a/mcp/schema_cache_test.go b/mcp/schema_cache_test.go new file mode 100644 index 0000000..b4d6ef6 --- /dev/null +++ b/mcp/schema_cache_test.go @@ -0,0 +1,208 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "context" + "reflect" + "testing" + + "github.com/google/jsonschema-go/jsonschema" +) + +func TestSchemaCacheByType(t *testing.T) { + cache := NewSchemaCache() + + type TestInput struct { + Name string `json:"name"` + } + + rt := reflect.TypeFor[TestInput]() + + if _, _, ok := cache.getByType(rt); ok { + t.Error("expected cache miss for new type") + } + + schema := &jsonschema.Schema{Type: "object"} + resolved, err := schema.Resolve(nil) + if err != nil { + t.Fatalf("failed to resolve schema: %v", err) + } + cache.setByType(rt, schema, resolved) + + gotSchema, gotResolved, ok := cache.getByType(rt) + if !ok { + t.Error("expected cache hit after set") + } + if gotSchema != schema { + t.Error("schema mismatch") + } + if gotResolved != resolved { + t.Error("resolved schema mismatch") + } +} + +func TestSchemaCacheBySchema(t *testing.T) { + cache := NewSchemaCache() + + schema := &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "query": {Type: "string"}, + }, + } + + if _, ok := cache.getBySchema(schema); ok { + t.Error("expected cache miss for new schema") + } + + resolved, err := schema.Resolve(nil) + if err != nil { + t.Fatalf("failed to resolve schema: %v", err) + } + cache.setBySchema(schema, resolved) + + gotResolved, ok := cache.getBySchema(schema) + if !ok { + t.Error("expected cache hit after set") + } + if gotResolved != resolved { + t.Error("resolved schema mismatch") + } + + // Different pointer should miss (cache uses pointer identity). + schema2 := &jsonschema.Schema{Type: "object"} + if _, ok = cache.getBySchema(schema2); ok { + t.Error("expected cache miss for different schema pointer") + } +} + +func TestSetSchemaCachesGeneratedSchemas(t *testing.T) { + cache := NewSchemaCache() + + type TestInput struct { + Query string `json:"query"` + } + + rt := reflect.TypeFor[TestInput]() + + var sfield1 any + var rfield1 *jsonschema.Resolved + if _, err := setSchema[TestInput](&sfield1, &rfield1, cache); err != nil { + t.Fatalf("setSchema failed: %v", err) + } + + cachedSchema, cachedResolved, ok := cache.getByType(rt) + if !ok { + t.Fatal("schema not cached after first setSchema call") + } + + var sfield2 any + var rfield2 *jsonschema.Resolved + if _, err := setSchema[TestInput](&sfield2, &rfield2, cache); err != nil { + t.Fatalf("setSchema failed on second call: %v", err) + } + + if sfield2.(*jsonschema.Schema) != cachedSchema { + t.Error("expected cached schema to be returned") + } + if rfield2 != cachedResolved { + t.Error("expected cached resolved schema to be returned") + } +} + +func TestSetSchemaCachesProvidedSchemas(t *testing.T) { + cache := NewSchemaCache() + + schema := &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "query": {Type: "string"}, + }, + } + + var sfield1 any = schema + var rfield1 *jsonschema.Resolved + if _, err := setSchema[map[string]any](&sfield1, &rfield1, cache); err != nil { + t.Fatalf("setSchema failed: %v", err) + } + + cachedResolved, ok := cache.getBySchema(schema) + if !ok { + t.Fatal("resolved schema not cached after first setSchema call") + } + if rfield1 != cachedResolved { + t.Error("expected same resolved schema") + } + + var sfield2 any = schema + var rfield2 *jsonschema.Resolved + if _, err := setSchema[map[string]any](&sfield2, &rfield2, cache); err != nil { + t.Fatalf("setSchema failed on second call: %v", err) + } + + if rfield2 != cachedResolved { + t.Error("expected cached resolved schema to be returned") + } +} + +func TestSetSchemaNilCache(t *testing.T) { + type TestInput struct { + Query string `json:"query"` + } + + var sfield1 any + var rfield1 *jsonschema.Resolved + if _, err := setSchema[TestInput](&sfield1, &rfield1, nil); err != nil { + t.Fatalf("setSchema failed: %v", err) + } + + var sfield2 any + var rfield2 *jsonschema.Resolved + if _, err := setSchema[TestInput](&sfield2, &rfield2, nil); err != nil { + t.Fatalf("setSchema failed on second call: %v", err) + } + + if sfield1 == nil || sfield2 == nil { + t.Error("expected schemas to be generated") + } + if rfield1 == nil || rfield2 == nil { + t.Error("expected resolved schemas to be generated") + } +} + +func TestAddToolWithSharedCache(t *testing.T) { + cache := NewSchemaCache() + + type GreetInput struct { + Name string `json:"name" jsonschema:"the name to greet"` + } + + type GreetOutput struct { + Message string `json:"message"` + } + + handler := func(ctx context.Context, req *CallToolRequest, in GreetInput) (*CallToolResult, GreetOutput, error) { + return &CallToolResult{}, GreetOutput{Message: "Hello, " + in.Name}, nil + } + + tool := &Tool{ + Name: "greet", + Description: "Greet someone", + } + + // Simulate stateless server pattern: new server per request, shared cache. + for range 3 { + s := NewServer(&Implementation{Name: "test", Version: "1.0"}, &ServerOptions{ + SchemaCache: cache, + }) + AddTool(s, tool, handler) + } + + rt := reflect.TypeFor[GreetInput]() + if _, _, ok := cache.getByType(rt); !ok { + t.Error("expected schema to be cached by type after multiple AddTool calls") + } +} diff --git a/mcp/server.go b/mcp/server.go index 29be8ff..569abab 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -7,9 +7,11 @@ package mcp import ( "bytes" "context" + "crypto/rand" "encoding/base64" "encoding/gob" "encoding/json" + "errors" "fmt" "iter" "log/slog" @@ -23,7 +25,9 @@ import ( "time" "github.com/google/jsonschema-go/jsonschema" + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/internal/mcpgodebug" "github.com/modelcontextprotocol/go-sdk/internal/util" "github.com/modelcontextprotocol/go-sdk/jsonrpc" "github.com/yosida95/uritemplate/v3" @@ -40,23 +44,36 @@ type Server struct { // fixed at creation impl *Implementation opts ServerOptions + // protocolVersions is the descending-ordered list of protocol versions + // this server advertises and negotiates: [supportedProtocolVersions], + // narrowed by [ServerOptions.SupportedProtocolVersions]. + protocolVersions []string - mu sync.Mutex - prompts *featureSet[*serverPrompt] - tools *featureSet[*serverTool] - resources *featureSet[*serverResource] - resourceTemplates *featureSet[*serverResourceTemplate] - sessions []*ServerSession - sendingMethodHandler_ MethodHandler - receivingMethodHandler_ MethodHandler - resourceSubscriptions map[string]map[*ServerSession]bool // uri -> session -> bool + mu sync.Mutex + prompts *featureSet[*serverPrompt] + tools *featureSet[*serverTool] + resources *featureSet[*serverResource] + resourceTemplates *featureSet[*serverResourceTemplate] + sessions []*ServerSession + sendingMethodHandler_ MethodHandler + receivingMethodHandler_ MethodHandler + toolChangeSubscriptions map[*ServerSession]jsonrpc.ID // session -> requestID for "tools/changed" + promptChangeSubscriptions map[*ServerSession]jsonrpc.ID // session -> requestID for "prompts/changed" + resourceChangeSubscriptions map[*ServerSession]jsonrpc.ID // session -> requestID for "resources/changed" + resourceSubscriptions map[string]map[*ServerSession]jsonrpc.ID // uri -> session -> requestID + pendingNotifications map[string]*time.Timer // notification name -> timer for pending notification send + // receiveMethods is the merged map of methods this server may receive + // from a client: it always contains the standard server methods (from + // serverMethodInfos) plus any custom methods registered via + // [AddReceivingCustomMethod]. + receiveMethods map[string]methodInfo } // ServerOptions is used to configure behavior of the server. type ServerOptions struct { // Optional instructions for connected clients. Instructions string - // If non-nil, log server activity. + // Logger may be set to a non-nil value to enable logging of server activity. Logger *slog.Logger // If non-nil, called when "notifications/initialized" is received. InitializedHandler func(context.Context, *InitializedRequest) @@ -66,6 +83,12 @@ type ServerOptions struct { // If zero, defaults to [DefaultPageSize]. PageSize int // If non-nil, called when "notifications/roots/list_changed" is received. + // + // Deprecated: the roots feature is deprecated as of protocol version + // 2026-07-28 (SEP-2577). It remains functional during the deprecation + // window (at least twelve months). Migrate to passing paths via tool + // parameters, resource URIs, or configuration. See + // https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. RootsListChangedHandler func(context.Context, *RootsListChangedRequest) // If non-nil, called when "notifications/progress" is received. ProgressNotificationHandler func(context.Context, *ProgressNotificationServerRequest) @@ -75,19 +98,68 @@ type ServerOptions struct { // If the peer fails to respond to pings originating from the keepalive check, // the session is automatically closed. KeepAlive time.Duration + // KeepAliveFailureThreshold is the number of consecutive keepalive ping + // failures tolerated before the session is closed. A value of 0 or 1 + // closes the session on the first failure (the default). Higher values + // align with the spec's "multiple failed pings MAY trigger a connection + // reset" guidance, letting a transient miss pass without tearing down an + // otherwise live session. Has no effect unless KeepAlive is non-zero. + KeepAliveFailureThreshold int // Function called when a client session subscribes to a resource. SubscribeHandler func(context.Context, *SubscribeRequest) error // Function called when a client session unsubscribes from a resource. UnsubscribeHandler func(context.Context, *UnsubscribeRequest) error + + // Capabilities optionally configures the server's default capabilities, + // before any capabilities are inferred from other configuration or server + // features. + // + // If Capabilities is nil, the default server capabilities are {"logging":{}}, + // for historical reasons. Setting Capabilities to a non-nil value overrides + // this default. For example, setting Capabilities to `&ServerCapabilities{}` + // disables the logging capability. + // + // # Interaction with capability inference + // + // "tools", "prompts", and "resources" capabilities are automatically added when + // tools, prompts, or resources are added to the server (for example, via + // [Server.AddPrompt]), with default value `{"listChanged":true}`. Similarly, + // if the [ClientOptions.SubscribeHandler] or + // [ClientOptions.CompletionHandler] are set, the inferred capabilities are + // adjusted accordingly. + // + // Any non-nil field in Capabilities overrides the inferred value. + // For example: + // + // - To advertise the "tools" capability, even if no tools are added, set + // Capabilities.Tools to &ToolCapabilities{ListChanged:true}. + // - To disable tool list notifications, set Capabilities.Tools to + // &ToolCapabilities{}. + // + // Conversely, if Capabilities does not set a field (for example, if the + // Prompts field is nil), the inferred capability will be used. + Capabilities *ServerCapabilities + // If true, advertises the prompts capability during initialization, // even if no prompts have been registered. + // + // Deprecated: Use Capabilities instead. HasPrompts bool // If true, advertises the resources capability during initialization, // even if no resources have been registered. + // + // Deprecated: Use Capabilities instead. HasResources bool // If true, advertises the tools capability during initialization, // even if no tools have been registered. + // + // Deprecated: Use Capabilities instead. HasTools bool + // SchemaCache, if non-nil, caches JSON schemas to avoid repeated + // reflection. This is useful for stateless server deployments where + // a new [Server] is created for each request. See [SchemaCache] for + // trade-offs and usage guidance. + SchemaCache *SchemaCache // GetSessionID provides the next session ID to use for an incoming request. // If nil, a default randomly generated ID will be used. @@ -97,7 +169,35 @@ type ServerOptions struct { // // As a special case, if GetSessionID returns the empty string, the // Mcp-Session-Id header will not be set. + // + // GetSessionID is not consulted when [StreamableHTTPOptions.Stateless] is + // true, since stateless servers do not maintain sessions. GetSessionID func() string + + // SetCacheable, if non-nil, decides the cache-control fields (ttlMs and + // cacheScope) of every result that carries them: server/discover, the + // four list methods, and resources/read. + // + // It is called once per such result, after the corresponding handler has + // returned, with c holding the values that handler produced. + SetCacheable func(ctx context.Context, req Request, c *Cacheable) + + // SupportedProtocolVersions, if non-empty, restricts the protocol versions + // this server advertises. If empty, every version returned by + // [SupportedProtocolVersions] is used. + // The list can only narrow support, never widen it. + // A request using the >= 2026-07-28 protocol at an excluded version is + // rejected with error code [CodeUnsupportedProtocolVersion]. + // + // The legacy initialize handshake never rejects. The lifecycle spec + // requires the server to answer a request it does not support with a + // version it does, so an excluded version is answered with the newest + // listed version that handshake can negotiate, and the client is expected + // to disconnect when it cannot speak it. A list holding no such version + // (one restricted to 2026-07-28 and later) is answered with 2025-11-25, + // the newest handshake-era version, so that the client disconnects rather + // than reading the answer as a negotiation into the new protocol. + SupportedProtocolVersions []string } // NewServer creates a new MCP server. The resulting server has no features: @@ -131,24 +231,47 @@ func NewServer(impl *Implementation, options *ServerOptions) *Server { } if opts.GetSessionID == nil { - opts.GetSessionID = randText + opts.GetSessionID = rand.Text + } + + protocolVersions := SupportedProtocolVersions() + if len(opts.SupportedProtocolVersions) > 0 { + for _, v := range opts.SupportedProtocolVersions { + if !slices.Contains(supportedProtocolVersions, v) { + panic(fmt.Errorf("unsupported protocol version %q", v)) + } + } + protocolVersions = slices.DeleteFunc(protocolVersions, func(v string) bool { + return !slices.Contains(opts.SupportedProtocolVersions, v) + }) } if opts.Logger == nil { // ensure we have a logger opts.Logger = ensureLogger(nil) } - return &Server{ - impl: impl, - opts: opts, - prompts: newFeatureSet(func(p *serverPrompt) string { return p.prompt.Name }), - tools: newFeatureSet(func(t *serverTool) string { return t.tool.Name }), - resources: newFeatureSet(func(r *serverResource) string { return r.resource.URI }), - resourceTemplates: newFeatureSet(func(t *serverResourceTemplate) string { return t.resourceTemplate.URITemplate }), - sendingMethodHandler_: defaultSendingMethodHandler[*ServerSession], - receivingMethodHandler_: defaultReceivingMethodHandler[*ServerSession], - resourceSubscriptions: make(map[string]map[*ServerSession]bool), + receiveMethods := make(map[string]methodInfo, len(serverMethodInfos)) + maps.Copy(receiveMethods, serverMethodInfos) + + s := &Server{ + impl: impl, + opts: opts, + protocolVersions: protocolVersions, + prompts: newFeatureSet(func(p *serverPrompt) string { return p.prompt.Name }), + tools: newFeatureSet(func(t *serverTool) string { return t.tool.Name }), + resources: newFeatureSet(func(r *serverResource) string { return r.resource.URI }), + resourceTemplates: newFeatureSet(func(t *serverResourceTemplate) string { return t.resourceTemplate.URITemplate }), + sendingMethodHandler_: defaultSendingMethodHandler, + receivingMethodHandler_: defaultReceivingMethodHandler[*ServerSession], + toolChangeSubscriptions: make(map[*ServerSession]jsonrpc.ID), + promptChangeSubscriptions: make(map[*ServerSession]jsonrpc.ID), + resourceChangeSubscriptions: make(map[*ServerSession]jsonrpc.ID), + resourceSubscriptions: make(map[string]map[*ServerSession]jsonrpc.ID), + pendingNotifications: make(map[string]*time.Timer), + receiveMethods: receiveMethods, } + s.AddReceivingMiddleware(serverMultiRoundTripMiddleware()) + return s } // AddPrompt adds a [Prompt] to the server, or replaces one with the same name. @@ -157,15 +280,13 @@ func (s *Server) AddPrompt(p *Prompt, h PromptHandler) { // (It's possible an item was replaced with an identical one, but not worth checking.) s.changeAndNotify( notificationPromptListChanged, - &PromptListChangedParams{}, func() bool { s.prompts.add(&serverPrompt{p, h}); return true }) } // RemovePrompts removes the prompts with the given names. // It is not an error to remove a nonexistent prompt. func (s *Server) RemovePrompts(names ...string) { - s.changeAndNotify(notificationPromptListChanged, &PromptListChangedParams{}, - func() bool { return s.prompts.remove(names...) }) + s.changeAndNotify(notificationPromptListChanged, func() bool { return s.prompts.remove(names...) }) } // AddTool adds a [Tool] to the server, or replaces one with the same name. @@ -175,7 +296,8 @@ func (s *Server) RemovePrompts(names ...string) { // that takes no input, or one where any input is valid, set [Tool.InputSchema] to // `{"type": "object"}`, using your preferred library or `json.RawMessage`. // -// If present, [Tool.OutputSchema] must also have type "object". +// If present, [Tool.OutputSchema] may be any valid JSON Schema (object, array, +// primitive, or composition). // // When the handler is invoked as part of a CallTool request, req.Params.Arguments // will be a json.RawMessage. @@ -191,6 +313,9 @@ func (s *Server) RemovePrompts(names ...string) { // Most users should use the top-level function [AddTool], which handles all these // responsibilities. func (s *Server) AddTool(t *Tool, h ToolHandler) { + if err := validateToolName(t.Name); err != nil { + s.opts.Logger.Error(fmt.Sprintf("AddTool: invalid tool name %q: %v", t.Name, err)) + } if t.InputSchema == nil { // This prevents the tool author from forgetting to write a schema where // one should be provided. If we papered over this by supplying the empty @@ -199,6 +324,9 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) { panic(fmt.Errorf("AddTool %q: missing input schema", t.Name)) } if s, ok := t.InputSchema.(*jsonschema.Schema); ok { + if s == nil { + panic(fmt.Errorf("AddTool %q: input schema is nil", t.Name)) + } if s.Type != "object" { panic(fmt.Errorf(`AddTool %q: input schema must have type "object"`, t.Name)) } @@ -213,29 +341,28 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) { } if t.OutputSchema != nil { if s, ok := t.OutputSchema.(*jsonschema.Schema); ok { - if s.Type != "object" { - panic(fmt.Errorf(`AddTool %q: output schema must have type "object"`, t.Name)) + if s == nil { + panic(fmt.Errorf("AddTool %q: output schema is nil", t.Name)) } } else { - var m map[string]any + var m any if err := remarshal(t.OutputSchema, &m); err != nil { - panic(fmt.Errorf("AddTool %q: can't marshal output schema to a JSON object: %v", t.Name, err)) - } - if typ := m["type"]; typ != "object" { - panic(fmt.Errorf(`AddTool %q: output schema must have type "object" (got %v)`, t.Name, typ)) + panic(fmt.Errorf("AddTool %q: can't marshal output schema to JSON: %v", t.Name, err)) } } } + if err := validateParamHeaderAnnotations(t); err != nil { + panic(fmt.Errorf("AddTool %q: invalid parameter header annotations: %v", t.Name, err)) + } st := &serverTool{tool: t, handler: h} // Assume there was a change, since add replaces existing tools. // (It's possible a tool was replaced with an identical one, but not worth checking.) // TODO: Batch these changes by size and time? The typescript SDK doesn't. // TODO: Surface notify error here? best not, in case we need to batch. - s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{}, - func() bool { s.tools.add(st); return true }) + s.changeAndNotify(notificationToolListChanged, func() bool { s.tools.add(st); return true }) } -func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHandler, error) { +func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out], cache *SchemaCache) (*Tool, ToolHandler, error) { tt := *t // Special handling for an "any" input: treat as an empty object. @@ -244,7 +371,7 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan } var inputResolved *jsonschema.Resolved - if _, err := setSchema[In](&tt.InputSchema, &inputResolved); err != nil { + if _, err := setSchema[In](&tt.InputSchema, &inputResolved, cache); err != nil { return nil, nil, fmt.Errorf("input schema: %w", err) } @@ -259,7 +386,7 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan ) if t.OutputSchema != nil || reflect.TypeFor[Out]() != reflect.TypeFor[any]() { var err error - elemZero, err = setSchema[Out](&tt.OutputSchema, &outputResolved) + elemZero, err = setSchema[Out](&tt.OutputSchema, &outputResolved, cache) if err != nil { return nil, nil, fmt.Errorf("output schema: %v", err) } @@ -272,34 +399,37 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan } // Validate input and apply defaults. var err error - input, err = applySchema(input, inputResolved) + input, err = applySchema(input, inputResolved, false) if err != nil { - // TODO(#450): should this be considered a tool error? (and similar below) - return nil, fmt.Errorf("%w: validating \"arguments\": %v", jsonrpc2.ErrInvalidParams, err) + var errRes CallToolResult + errRes.SetError(fmt.Errorf("validating \"arguments\": %v", err)) + return &errRes, nil } // Unmarshal and validate args. var in In if input != nil { - if err := json.Unmarshal(input, &in); err != nil { - return nil, fmt.Errorf("%w: %v", jsonrpc2.ErrInvalidParams, err) + if err := internaljson.Unmarshal(input, &in); err != nil { + var errRes CallToolResult + errRes.SetError(err) + return &errRes, nil } } // Call typed handler. res, out, err := h(ctx, req, in) // Handle server errors appropriately: - // - If the handler returns a structured error (like jsonrpc2.WireError), return it directly + // - If the handler returns a structured error (like jsonrpc.Error), return it directly // - If the handler returns a regular error, wrap it in a CallToolResult with IsError=true // - This allows tools to distinguish between protocol errors and tool execution errors if err != nil { // Check if this is already a structured JSON-RPC error - if wireErr, ok := err.(*jsonrpc2.WireError); ok { + if wireErr, ok := err.(*jsonrpc.Error); ok { return nil, wireErr } // For regular errors, embed them in the tool result as per MCP spec var errRes CallToolResult - errRes.setError(err) + errRes.SetError(err) return &errRes, nil } @@ -308,8 +438,12 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan } // Marshal the output and put the RawMessage in the StructuredContent field. + // Skip when the handler returned input requests (multi round-trip): content and + // inputRequests are mutually exclusive on the wire. var outval any = out - if elemZero != nil { + if res.InputRequests != nil { + outval = nil + } else if elemZero != nil { // Avoid typed nil, which will serialize as JSON null. // Instead, use the zero value of the unpointered type. var z Out @@ -327,7 +461,7 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan // // We validate against the JSON, rather than the output value, as // some types may have custom JSON marshalling (issue #447). - outJSON, err = applySchema(outJSON, outputResolved) + outJSON, err = applySchema(outJSON, outputResolved, true) if err != nil { return nil, fmt.Errorf("validating tool output: %w", err) } @@ -336,10 +470,18 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan // If the Content field isn't being used, return the serialized JSON in a // TextContent block, as the spec suggests: // https://modelcontextprotocol.io/specification/2025-06-18/server/tools#structured-content. + // + // Ensure a serialized-JSON TextContent fallback is present in case of servers using array + // or primitive structuredContent, so that pre-SEP-2106 clients can recover the structured + // payload from unstructured content. if res.Content == nil { res.Content = []Content{&TextContent{ Text: string(outJSON), }} + } else if !isObjectJSON(outJSON) { + res.Content = append(res.Content, &TextContent{ + Text: string(outJSON), + }) } } return res, nil @@ -360,29 +502,82 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan // pointer: if the user provided the schema, they may have intentionally // derived it from the pointer type, and handling of zero values is up to them. // +// If cache is non-nil, schemas are cached to avoid repeated reflection. +// // TODO(rfindley): we really shouldn't ever return 'null' results. Maybe we // should have a jsonschema.Zero(schema) helper? -func setSchema[T any](sfield *any, rfield **jsonschema.Resolved) (zero any, err error) { - var internalSchema *jsonschema.Schema - if *sfield == nil { - rt := reflect.TypeFor[T]() - if rt.Kind() == reflect.Pointer { - rt = rt.Elem() - zero = reflect.Zero(rt).Interface() - } - // TODO: we should be able to pass nil opts here. - internalSchema, err = jsonschema.ForType(rt, &jsonschema.ForOptions{}) - if err == nil { - *sfield = internalSchema - } - } else if err := remarshal(*sfield, &internalSchema); err != nil { - return zero, err +func setSchema[T any](sfield *any, rfield **jsonschema.Resolved, cache *SchemaCache) (zero any, err error) { + rt := reflect.TypeFor[T]() + if rt.Kind() == reflect.Pointer { + rt = rt.Elem() + zero = reflect.Zero(rt).Interface() } + + var internalSchema *jsonschema.Schema + + if *sfield == nil { + // No schema provided: check cache, or generate via reflection. + if cache != nil { + if schema, resolved, ok := cache.getByType(rt); ok { + *sfield = schema + *rfield = resolved + return zero, nil + } + } + + internalSchema, err = jsonschema.ForType(rt, &jsonschema.ForOptions{}) + if err != nil { + return zero, err + } + if internalSchema == nil { + return zero, fmt.Errorf("schema is nil for type %v", rt) + } + *sfield = internalSchema + + resolved, err := internalSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true}) + if err != nil { + return zero, err + } + *rfield = resolved + if cache != nil { + cache.setByType(rt, internalSchema, resolved) + } + return zero, nil + } + + // Schema was provided: check cache by pointer, or resolve it. + if providedSchema, ok := (*sfield).(*jsonschema.Schema); ok { + if cache != nil { + if resolved, ok := cache.getBySchema(providedSchema); ok { + *rfield = resolved + return zero, nil + } + } + internalSchema = providedSchema + } else { + // Schema provided as different type (e.g., map): remarshal to *Schema. + if err := remarshal(*sfield, &internalSchema); err != nil { + return zero, err + } + } + + if internalSchema == nil { + return zero, fmt.Errorf("schema is nil for type %v", rt) + } + + resolved, err := internalSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true}) if err != nil { return zero, err } - *rfield, err = internalSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true}) - return zero, err + *rfield = resolved + + if cache != nil { + if providedSchema, ok := (*sfield).(*jsonschema.Schema); ok { + cache.setBySchema(providedSchema, resolved) + } + } + + return zero, nil } // AddTool adds a tool and typed tool handler to the server. @@ -397,15 +592,16 @@ func setSchema[T any](sfield *any, rfield **jsonschema.Resolved) (zero any, err // empty object schema value. // // If the tool's output schema is nil, and the Out type is not 'any', the -// output schema is set to the schema inferred from the Out type argument, -// which must also be a map or struct. If the Out type is 'any', the output -// schema is omitted. +// output schema is set to the schema inferred from the Out type argument. +// Per SEP-2106, the Out type may be any Go type whose inferred schema is a +// valid JSON Schema (struct, map, slice, primitive, etc.). If the Out type is +// 'any', the output schema is omitted. // // Unlike [Server.AddTool], AddTool does a lot automatically, and forces // tools to conform to the MCP spec. See [ToolHandlerFor] for a detailed // description of this automatic behavior. func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) { - tt, hh, err := toolForErr(t, h) + tt, hh, err := toolForErr(t, h, s.opts.SchemaCache) if err != nil { panic(fmt.Sprintf("AddTool: tool %q: %v", t.Name, err)) } @@ -415,14 +611,13 @@ func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) { // RemoveTools removes the tools with the given names. // It is not an error to remove a nonexistent tool. func (s *Server) RemoveTools(names ...string) { - s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{}, - func() bool { return s.tools.remove(names...) }) + s.changeAndNotify(notificationToolListChanged, func() bool { return s.tools.remove(names...) }) } // AddResource adds a [Resource] to the server, or replaces one with the same URI. // AddResource panics if the resource URI is invalid or not absolute (has an empty scheme). func (s *Server) AddResource(r *Resource, h ResourceHandler) { - s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{}, + s.changeAndNotify(notificationResourceListChanged, func() bool { if _, err := url.Parse(r.URI); err != nil { panic(err) // url.Parse includes the URI in the error @@ -435,14 +630,13 @@ func (s *Server) AddResource(r *Resource, h ResourceHandler) { // RemoveResources removes the resources with the given URIs. // It is not an error to remove a nonexistent resource. func (s *Server) RemoveResources(uris ...string) { - s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{}, - func() bool { return s.resources.remove(uris...) }) + s.changeAndNotify(notificationResourceListChanged, func() bool { return s.resources.remove(uris...) }) } // AddResourceTemplate adds a [ResourceTemplate] to the server, or replaces one with the same URI. // AddResourceTemplate panics if a URI template is invalid or not absolute (has an empty scheme). func (s *Server) AddResourceTemplate(t *ResourceTemplate, h ResourceHandler) { - s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{}, + s.changeAndNotify(notificationResourceListChanged, func() bool { // Validate the URI template syntax _, err := uritemplate.New(t.URITemplate) @@ -457,54 +651,220 @@ func (s *Server) AddResourceTemplate(t *ResourceTemplate, h ResourceHandler) { // RemoveResourceTemplates removes the resource templates with the given URI templates. // It is not an error to remove a nonexistent resource. func (s *Server) RemoveResourceTemplates(uriTemplates ...string) { - s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{}, - func() bool { return s.resourceTemplates.remove(uriTemplates...) }) + s.changeAndNotify(notificationResourceListChanged, func() bool { return s.resourceTemplates.remove(uriTemplates...) }) } func (s *Server) capabilities() *ServerCapabilities { s.mu.Lock() defer s.mu.Unlock() - caps := &ServerCapabilities{ - Logging: &LoggingCapabilities{}, + // Start with user-provided capabilities as defaults, or use SDK defaults. + var caps *ServerCapabilities + if s.opts.Capabilities != nil { + // Deep copy the user-provided capabilities to avoid mutation. + caps = s.opts.Capabilities.clone() + } else { + // SDK defaults: only logging capability. + caps = &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + } } + + // Augment with tools capability if tools exist or legacy HasTools is set. if s.opts.HasTools || s.tools.len() > 0 { - caps.Tools = &ToolCapabilities{ListChanged: true} + if caps.Tools == nil { + caps.Tools = &ToolCapabilities{ListChanged: true} + } } + + // Augment with prompts capability if prompts exist or legacy HasPrompts is set. if s.opts.HasPrompts || s.prompts.len() > 0 { - caps.Prompts = &PromptCapabilities{ListChanged: true} + if caps.Prompts == nil { + caps.Prompts = &PromptCapabilities{ListChanged: true} + } } + + // Augment with resources capability if resources/templates exist or legacy HasResources is set. if s.opts.HasResources || s.resources.len() > 0 || s.resourceTemplates.len() > 0 { - caps.Resources = &ResourceCapabilities{ListChanged: true} + if caps.Resources == nil { + caps.Resources = &ResourceCapabilities{ListChanged: true} + } if s.opts.SubscribeHandler != nil { caps.Resources.Subscribe = true } } + + // Augment with completions capability if handler is set. if s.opts.CompletionHandler != nil { - caps.Completions = &CompletionCapabilities{} + if caps.Completions == nil { + caps.Completions = &CompletionCapabilities{} + } } + return caps } +// disablecompleteparamsvalidation is a compatibility parameter that restores +// the previous behavior of [Server.complete], where required fields on +// [CompleteParams] were not validated before dispatching to the completion +// handler. See the documentation for the mcpgodebug package for instructions +// how to enable it. +// The option will be removed in a future version of the SDK. +var disablecompleteparamsvalidation = mcpgodebug.Value("disablecompleteparamsvalidation") + func (s *Server) complete(ctx context.Context, req *CompleteRequest) (*CompleteResult, error) { + if disablecompleteparamsvalidation != "1" { + if req.Params.Ref == nil { + return nil, fmt.Errorf("%w: missing required 'ref' field", jsonrpc2.ErrInvalidParams) + } + if req.Params.Argument.Name == "" { + return nil, fmt.Errorf("%w: missing required 'argument.name' field", jsonrpc2.ErrInvalidParams) + } + } if s.opts.CompletionHandler == nil { return nil, jsonrpc2.ErrMethodNotFound } return s.opts.CompletionHandler(ctx, req) } +// Map from notification name to a function creating its corresponding Params. +// We need to create a fresh one each time to add the jsonrpc ID. See +// [injectMetaSubscriptionID]. +var changeNotificationParams = map[string]func() Params{ + notificationToolListChanged: func() Params { return &ToolListChangedParams{} }, + notificationPromptListChanged: func() Params { return &PromptListChangedParams{} }, + notificationResourceListChanged: func() Params { return &ResourceListChangedParams{} }, +} + +// How long to wait before sending a change notification. +const notificationDelay = 10 * time.Millisecond + // changeAndNotify is called when a feature is added or removed. // It calls change, which should do the work and report whether a change actually occurred. -// If there was a change, it notifies a snapshot of the sessions. -func (s *Server) changeAndNotify(notification string, params Params, change func() bool) { - var sessions []*ServerSession - // Lock for the change, but not for the notification. +// If there was a change, it sets a timer to send a notification. +// This debounces change notifications: a single notification is sent after +// multiple changes occur in close proximity. +func (s *Server) changeAndNotify(notification string, change func() bool) { s.mu.Lock() - if change() { - sessions = slices.Clone(s.sessions) + defer s.mu.Unlock() + if change() && s.shouldSendListChangedNotification(notification) { + if len(s.sessions) == 0 { + if t := s.pendingNotifications[notification]; t != nil { + t.Stop() + s.pendingNotifications[notification] = nil + } + return + } + + // Reset the outstanding delayed call, if any. + if t := s.pendingNotifications[notification]; t == nil { + s.pendingNotifications[notification] = time.AfterFunc(notificationDelay, func() { s.notifySessions(notification) }) + } else { + t.Reset(notificationDelay) + } + } +} + +// notifySessions sends the notification n to all existing sessions. +// It is called asynchronously by changeAndNotify. +// +// Legacy (pre-SEP-2575) sessions receive the notification on the shared +// session channel. Sessions speaking the new protocol receive it only if they +// have an active subscriptions/listen stream that opted in to this +// notification type. +func (s *Server) notifySessions(n string) { + s.mu.Lock() + s.pendingNotifications[n] = nil + // Legacy (pre-SEP-2575) sessions receive list-changed notifications on the + // shared session channel without opt-in; collect them while we hold the lock. + var legacySessions []*ServerSession + for _, sess := range s.sessions { + if sess.InitializeParams().isNil() || sess.InitializeParams().ProtocolVersion < protocolVersion20260728 { + legacySessions = append(legacySessions, sess) + } + } + var subscribers map[*ServerSession]jsonrpc.ID + switch n { + case notificationToolListChanged: + subscribers = maps.Clone(s.toolChangeSubscriptions) + case notificationPromptListChanged: + subscribers = maps.Clone(s.promptChangeSubscriptions) + case notificationResourceListChanged: + subscribers = maps.Clone(s.resourceChangeSubscriptions) + } + s.mu.Unlock() // Don't hold the lock during notification: it causes deadlock. + + // Notify legacy sessions on the shared session channel regardless of + // their subscription status. + notifySessions(legacySessions, n, changeNotificationParams[n](), s.opts.Logger) + + // Notify modern sessions only if they have an active subscription for + // this notification type. + s.notifySubscribedSessions(subscribers, n, changeNotificationParams[n]) +} + +// notifySubscribedSessions delivers a list-changed or resource-updated +// notification to each session in subscribers, stamping the session's +// listen-request ID into the params' _meta so the receiving client can demultiplex +// notifications belonging to different concurrent listens. +func (s *Server) notifySubscribedSessions(subscribers map[*ServerSession]jsonrpc.ID, method string, makeParams func() Params) { + if len(subscribers) == 0 { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + for sess, reqID := range subscribers { + params := makeParams() + injectMetaSubscriptionID(params, reqID) + req := newRequest(sess, params) + if err := handleNotify(ctx, method, req); err != nil { + s.opts.Logger.Warn(fmt.Sprintf("calling %s: %v", method, err)) + } + } +} + +// injectMetaSubscriptionID stamps the listen request's JSON-RPC ID into the +// params' _meta under [MetaKeySubscriptionID], so the receiving client can +// correlate the notification with the [subscriptions/listen] stream it +// originated. +// +// [subscriptions/listen]: https://modelcontextprotocol.io/seps/2575-stateless-mcp#multiple-concurrent-subscriptions +func injectMetaSubscriptionID(params Params, reqID jsonrpc.ID) { + m := params.GetMeta() + if m == nil { + m = map[string]any{} + } + m[MetaKeySubscriptionID] = reqID.Raw() + params.SetMeta(m) +} + +// shouldSendListChangedNotification checks if the server's capabilities allow +// sending the given list-changed notification. +func (s *Server) shouldSendListChangedNotification(notification string) bool { + // Get effective capabilities (considering user-provided defaults). + caps := s.opts.Capabilities + + switch notification { + case notificationToolListChanged: + // If user didn't specify capabilities, default behavior sends notifications. + if caps == nil || caps.Tools == nil { + return true + } + return caps.Tools.ListChanged + case notificationPromptListChanged: + if caps == nil || caps.Prompts == nil { + return true + } + return caps.Prompts.ListChanged + case notificationResourceListChanged: + if caps == nil || caps.Resources == nil { + return true + } + return caps.Resources.ListChanged + default: + // Unknown notification, allow by default. + return true } - s.mu.Unlock() - notifySessions(sessions, notification, params) } // Sessions returns an iterator that yields the current set of server sessions. @@ -518,18 +878,23 @@ func (s *Server) Sessions() iter.Seq[*ServerSession] { return slices.Values(clients) } -func (s *Server) listPrompts(_ context.Context, req *ListPromptsRequest) (*ListPromptsResult, error) { +func (s *Server) listPrompts(ctx context.Context, req *ListPromptsRequest) (*ListPromptsResult, error) { s.mu.Lock() defer s.mu.Unlock() if req.Params == nil { req.Params = &ListPromptsParams{} } - return paginateList(s.prompts, s.opts.PageSize, req.Params, &ListPromptsResult{}, func(res *ListPromptsResult, prompts []*serverPrompt) { + res, err := paginateList(s.prompts, s.opts.PageSize, req.Params, &ListPromptsResult{}, func(res *ListPromptsResult, prompts []*serverPrompt) { res.Prompts = []*Prompt{} // avoid JSON null for _, p := range prompts { res.Prompts = append(res.Prompts, p.prompt) } }) + if err != nil { + return nil, err + } + s.resolveCacheable(ctx, req, &res.Cacheable) + return res, nil } func (s *Server) getPrompt(ctx context.Context, req *GetPromptRequest) (*GetPromptResult, error) { @@ -538,74 +903,164 @@ func (s *Server) getPrompt(ctx context.Context, req *GetPromptRequest) (*GetProm s.mu.Unlock() if !ok { // Return a proper JSON-RPC error with the correct error code - return nil, &jsonrpc2.WireError{ - Code: codeInvalidParams, + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("unknown prompt %q", req.Params.Name), } } - return prompt.handler(ctx, req) + res, err := prompt.handler(ctx, req) + if err == nil && res != nil { + if err := validateMultiRoundTripResult(s.opts.Logger, res); err != nil { + return nil, err + } + } + return res, err } -func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListToolsResult, error) { +// discover is the server-side handler for the SEP-2575 "server/discover" RPC. +// +// It returns the protocol versions supported by the underlying transport, +// the server's capabilities, the server's identity, and the server's +// instructions, allowing clients to negotiate without performing the legacy +// initialize handshake. +func (s *Server) discover(ctx context.Context, req *ServerRequest[*DiscoverParams]) (*DiscoverResult, error) { + req.Session.mu.Lock() + versions := req.Session.supportedVersions + req.Session.mu.Unlock() + if versions == nil { + versions = slices.Clone(s.protocolVersions) + } + // Read the request-scoped identity/capabilities before acquiring the + // session lock: these accessors may fall back to Session.InitializeParams + // (which also locks Session.mu), so calling them from inside updateState + // would self-deadlock. + init := &InitializeParams{ + ProtocolVersion: req.ProtocolVersion(), + Capabilities: req.ClientCapabilities(), + ClientInfo: req.ClientInfo(), + } + // Only persist InitializeParams when the transport can actually serve + // the new protocol. On transports that cannot (notably stateful + // StreamableHTTPHandler), a discover request creates a session that + // is never surfaced to the client via Mcp-Session-Id; leaving + // InitializeParams nil lets serveStatefulPOST's safety-net cleanup + // close it instead of leaking. + if slices.ContainsFunc(versions, func(v string) bool { return v >= protocolVersion20260728 }) { + req.Session.updateState(func(state *ServerSessionState) { + state.InitializeParams = init + }) + } + res := &DiscoverResult{ + SupportedVersions: versions, + Capabilities: s.capabilities(), + Instructions: s.opts.Instructions, + } + s.resolveCacheable(ctx, req, &res.Cacheable) + return res, nil +} + +// filterSupportedVersions returns the subset of versions that the Transport +// can serve. If t does not implement [ProtocolVersionSupporter], every version +// is included. +func filterSupportedVersions(t Transport, versions []string) []string { + pvs, ok := t.(ProtocolVersionSupporter) + if !ok { + return slices.Clone(versions) + } + out := make([]string, 0, len(versions)) + for _, v := range versions { + if pvs.SupportsProtocolVersion(v) { + out = append(out, v) + } + } + return out +} + +func (s *Server) listTools(ctx context.Context, req *ListToolsRequest) (*ListToolsResult, error) { s.mu.Lock() defer s.mu.Unlock() if req.Params == nil { req.Params = &ListToolsParams{} } - return paginateList(s.tools, s.opts.PageSize, req.Params, &ListToolsResult{}, func(res *ListToolsResult, tools []*serverTool) { + res, err := paginateList(s.tools, s.opts.PageSize, req.Params, &ListToolsResult{}, func(res *ListToolsResult, tools []*serverTool) { res.Tools = []*Tool{} // avoid JSON null for _, t := range tools { res.Tools = append(res.Tools, t.tool) } }) + if err != nil { + return nil, err + } + s.resolveCacheable(ctx, req, &res.Cacheable) + return res, nil +} + +// getServerTool looks up a server tool by name. +func (s *Server) getServerTool(name string) (*serverTool, bool) { + s.mu.Lock() + defer s.mu.Unlock() + return s.tools.get(name) } func (s *Server) callTool(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { - s.mu.Lock() - st, ok := s.tools.get(req.Params.Name) - s.mu.Unlock() + st, ok := s.getServerTool(req.Params.Name) if !ok { - return nil, &jsonrpc2.WireError{ - Code: codeInvalidParams, + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("unknown tool %q", req.Params.Name), } } res, err := st.handler(ctx, req) - if err == nil && res != nil && res.Content == nil { - res2 := *res - res2.Content = []Content{} // avoid "null" - res = &res2 + if err == nil && res != nil { + if err := validateMultiRoundTripResult(s.opts.Logger, res); err != nil { + return nil, err + } + if res.Content == nil && res.InputRequests == nil { + res2 := *res + res2.Content = []Content{} // avoid "null" + res = &res2 + } } return res, err } -func (s *Server) listResources(_ context.Context, req *ListResourcesRequest) (*ListResourcesResult, error) { +func (s *Server) listResources(ctx context.Context, req *ListResourcesRequest) (*ListResourcesResult, error) { s.mu.Lock() defer s.mu.Unlock() if req.Params == nil { req.Params = &ListResourcesParams{} } - return paginateList(s.resources, s.opts.PageSize, req.Params, &ListResourcesResult{}, func(res *ListResourcesResult, resources []*serverResource) { + res, err := paginateList(s.resources, s.opts.PageSize, req.Params, &ListResourcesResult{}, func(res *ListResourcesResult, resources []*serverResource) { res.Resources = []*Resource{} // avoid JSON null for _, r := range resources { res.Resources = append(res.Resources, r.resource) } }) + if err != nil { + return nil, err + } + s.resolveCacheable(ctx, req, &res.Cacheable) + return res, nil } -func (s *Server) listResourceTemplates(_ context.Context, req *ListResourceTemplatesRequest) (*ListResourceTemplatesResult, error) { +func (s *Server) listResourceTemplates(ctx context.Context, req *ListResourceTemplatesRequest) (*ListResourceTemplatesResult, error) { s.mu.Lock() defer s.mu.Unlock() if req.Params == nil { req.Params = &ListResourceTemplatesParams{} } - return paginateList(s.resourceTemplates, s.opts.PageSize, req.Params, &ListResourceTemplatesResult{}, + res, err := paginateList(s.resourceTemplates, s.opts.PageSize, req.Params, &ListResourceTemplatesResult{}, func(res *ListResourceTemplatesResult, rts []*serverResourceTemplate) { res.ResourceTemplates = []*ResourceTemplate{} // avoid JSON null for _, rt := range rts { res.ResourceTemplates = append(res.ResourceTemplates, rt.resourceTemplate) } }) + if err != nil { + return nil, err + } + s.resolveCacheable(ctx, req, &res.Cacheable) + return res, nil } func (s *Server) readResource(ctx context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { @@ -622,7 +1077,17 @@ func (s *Server) readResource(ctx context.Context, req *ReadResourceRequest) (*R if err != nil { return nil, err } - if res == nil || res.Contents == nil { + if res == nil { + return nil, fmt.Errorf("reading resource %s: read handler returned nil information", uri) + } + if err := validateMultiRoundTripResult(s.opts.Logger, res); err != nil { + return nil, err + } + s.resolveCacheable(ctx, req, &res.Cacheable) + if res.InputRequests != nil { + return res, nil + } + if res.Contents == nil { return nil, fmt.Errorf("reading resource %s: read handler returned nil information", uri) } // As a convenience, populate some fields. @@ -637,6 +1102,22 @@ func (s *Server) readResource(ctx context.Context, req *ReadResourceRequest) (*R return res, nil } +// resolveCacheable settles the cache-control fields of an outgoing result. +// +// On entry c holds whatever the feature handler produced, which is the zero +// value for results the SDK builds itself. [ServerOptions.SetCacheable] may +// then rewrite it, and any cacheScope still unset is filled with the protocol +// default. +// +// The list methods call this with s.mu held, so SetCacheable must not call +// back into the server; see its documentation. +func (s *Server) resolveCacheable(ctx context.Context, req Request, c *Cacheable) { + if s.opts.SetCacheable != nil { + s.opts.SetCacheable(ctx, req, c) + } + c.normalize() +} + // lookupResourceHandler returns the resource handler and MIME type for the resource or // resource template matching uri. If none, the last return value is false. func (s *Server) lookupResourceHandler(uri string) (ResourceHandler, string, bool) { @@ -665,7 +1146,7 @@ func (s *Server) lookupResourceHandler(uri string) (ResourceHandler, string, boo // and the current working directory is unavailable, fileResourceHandler panics. // // Lexical path traversal attacks, where the path has ".." elements that escape dir, -// are always caught. Go 1.24 and above also protects against symlink-based attacks, +// are always caught. The SDK also protects against symlink-based attacks, // where symlinks under dir lead out of the tree. func fileResourceHandler(dir string) ResourceHandler { // Convert dir to an absolute path. @@ -702,14 +1183,36 @@ func fileResourceHandler(dir string) ResourceHandler { func (s *Server) ResourceUpdated(ctx context.Context, params *ResourceUpdatedNotificationParams) error { s.mu.Lock() subscribedSessions := s.resourceSubscriptions[params.URI] - sessions := slices.Collect(maps.Keys(subscribedSessions)) + // Only add legacy sessions for the notification, new ones use the new notification mechanism. + var legacySessions []*ServerSession + newSessions := make(map[*ServerSession]jsonrpc.ID) + for sess, reqID := range subscribedSessions { + if sess.InitializeParams().isNil() || sess.InitializeParams().ProtocolVersion < protocolVersion20260728 { + legacySessions = append(legacySessions, sess) + } else { + newSessions[sess] = reqID + } + } s.mu.Unlock() - notifySessions(sessions, notificationResourceUpdated, params) - s.opts.Logger.Info("resource updated notification sent", "uri", params.URI, "subscriber_count", len(sessions)) + + notifySessions(legacySessions, notificationResourceUpdated, params, s.opts.Logger) + + // Notify modern sessions, injecting the per-session subscription ID into the notification's metadata. + s.notifySubscribedSessions(newSessions, notificationResourceUpdated, func() Params { + p := *params + return &p + }) + s.opts.Logger.Info("resource updated notification sent", + "uri", params.URI, + "subscriber_count", len(legacySessions)+len(newSessions)) return nil } func (s *Server) subscribe(ctx context.Context, req *SubscribeRequest) (*emptyResult, error) { + requestID, ok := ctx.Value(idContextKey{}).(jsonrpc.ID) + if !ok || !requestID.IsValid() { + return nil, fmt.Errorf("%w: subscribe requires a request ID", jsonrpc2.ErrInvalidRequest) + } if s.opts.SubscribeHandler == nil { return nil, fmt.Errorf("%w: server does not support resource subscriptions", jsonrpc2.ErrMethodNotFound) } @@ -720,10 +1223,10 @@ func (s *Server) subscribe(ctx context.Context, req *SubscribeRequest) (*emptyRe s.mu.Lock() defer s.mu.Unlock() if s.resourceSubscriptions[req.Params.URI] == nil { - s.resourceSubscriptions[req.Params.URI] = make(map[*ServerSession]bool) + s.resourceSubscriptions[req.Params.URI] = make(map[*ServerSession]jsonrpc.ID) } - s.resourceSubscriptions[req.Params.URI][req.Session] = true - s.opts.Logger.Info("resource subscribed", "uri", req.Params.URI, "session_id", req.Session.ID()) + s.resourceSubscriptions[req.Params.URI][req.Session] = requestID + s.opts.Logger.Info("resource subscribed", "uri", req.Params.URI, "session_id", req.Session.ID(), "request_id", requestID) return &emptyResult{}, nil } @@ -750,6 +1253,100 @@ func (s *Server) unsubscribe(ctx context.Context, req *UnsubscribeRequest) (*emp return &emptyResult{}, nil } +func (s *Server) subscriptionsListen(ctx context.Context, req *SubscriptionsListenRequest) (*SubscriptionsListenResult, error) { + requestID, ok := ctx.Value(idContextKey{}).(jsonrpc.ID) + if !ok || !requestID.IsValid() { + return nil, fmt.Errorf("%w: subscriptions/listen requires a request ID", jsonrpc2.ErrInvalidRequest) + } + + if req.Params.Notifications == nil { + return nil, fmt.Errorf("%w: missing required 'notifications' field", jsonrpc2.ErrInvalidParams) + } + + allowed := s.allowedSubscriptions(req.Params.Notifications) + s.mu.Lock() + if allowed.ToolsListChanged { + s.toolChangeSubscriptions[req.Session] = requestID + } + if allowed.PromptsListChanged { + s.promptChangeSubscriptions[req.Session] = requestID + } + if allowed.ResourcesListChanged { + s.resourceChangeSubscriptions[req.Session] = requestID + } + s.mu.Unlock() + defer func() { + // Retire only what this listen registered above: a session can hold + // several listens at once, and deleting by session alone retires a + // registration belonging to a listen that is still open. + s.mu.Lock() + if allowed.ToolsListChanged { + delete(s.toolChangeSubscriptions, req.Session) + } + if allowed.PromptsListChanged { + delete(s.promptChangeSubscriptions, req.Session) + } + if allowed.ResourcesListChanged { + delete(s.resourceChangeSubscriptions, req.Session) + } + s.mu.Unlock() + }() + + for _, uri := range allowed.ResourceSubscriptions { + _, err := s.subscribe(ctx, &SubscribeRequest{ + Session: req.Session, + Params: &SubscribeParams{ + URI: uri, + Meta: req.Params.GetMeta(), + }, + }) + if err != nil { + return nil, err + } + defer s.unsubscribe(ctx, &UnsubscribeRequest{ + Session: req.Session, + Params: &UnsubscribeParams{ + URI: uri, + Meta: req.Params.GetMeta(), + }, + }) + } + + ackParams := &SubscriptionsAcknowledgedParams{ + Notifications: allowed, + Meta: Meta{MetaKeySubscriptionID: requestID.Raw()}, + } + if err := req.Session.notifySubscriptionAcked(ctx, ackParams); err != nil { + return nil, fmt.Errorf("sending subscriptions/acknowledged: %w", err) + } + + // If there are any active subscriptions, we block until the context is cancelled. Otherwise, we return immediately. + if len(allowed.ResourceSubscriptions) > 0 || allowed.ToolsListChanged || allowed.PromptsListChanged || allowed.ResourcesListChanged { + <-ctx.Done() + } + return &SubscriptionsListenResult{ + Meta: Meta{MetaKeySubscriptionID: requestID.Raw()}, + }, nil +} + +func (s *Server) allowedSubscriptions(want *NotificationSubscriptions) NotificationSubscriptions { + caps := s.capabilities() + agreed := NotificationSubscriptions{} + if want.ToolsListChanged && caps.Tools != nil && caps.Tools.ListChanged { + agreed.ToolsListChanged = true + } + if want.PromptsListChanged && caps.Prompts != nil && caps.Prompts.ListChanged { + agreed.PromptsListChanged = true + } + if want.ResourcesListChanged && caps.Resources != nil && caps.Resources.ListChanged { + agreed.ResourcesListChanged = true + } + if len(want.ResourceSubscriptions) > 0 && caps.Resources != nil && caps.Resources.Subscribe { + agreed.ResourceSubscriptions = slices.Clone(want.ResourceSubscriptions) + } + return agreed +} + // Run runs the server over the given transport, which must be persistent. // // Run blocks until the client terminates the connection or the provided @@ -819,6 +1416,10 @@ func (s *Server) disconnect(cc *ServerSession) { for _, subscribedSessions := range s.resourceSubscriptions { delete(subscribedSessions, cc) } + delete(s.toolChangeSubscriptions, cc) + delete(s.promptChangeSubscriptions, cc) + delete(s.resourceChangeSubscriptions, cc) + s.opts.Logger.Info("server session disconnected", "session_id", cc.ID()) } @@ -846,11 +1447,33 @@ func (s *Server) Connect(ctx context.Context, t Transport, opts *ServerSessionOp } s.opts.Logger.Info("server connecting") - ss, err := connect(ctx, t, s, state, onClose) + ss, err := connect(ctx, t, s, state, onClose, s.opts.Logger) if err != nil { s.opts.Logger.Error("server connect error", "error", err) return nil, err } + + // Compute the protocol versions this session can serve, filtered by the + // transport's capabilities (if it implements [ProtocolVersionSupporter]). + // The list is consumed by the SEP-2575 server/discover handler. + // + // The write is guarded by ss.mu to establish a happens-before edge with + // the matching read in Server.discover, which runs on the jsonrpc2 read + // goroutine spawned inside connect(). The two are not concurrent in + // wall-clock terms (no incoming message is dispatched until the caller + // has fed the transport, which happens after Server.Connect returns), + // but without the lock the Go memory model gives the read goroutine no + // guarantee of seeing this write, and -race flags it. + ss.mu.Lock() + ss.supportedVersions = filterSupportedVersions(t, s.protocolVersions) + ss.mu.Unlock() + + // Start keepalive before returning the session to avoid race conditions with Close. + // This is safe because the spec allows sending pings before initialization (see ServerSession.handle for details). + if s.opts.KeepAlive > 0 { + ss.startKeepalive(ss.server.opts.KeepAlive) + } + return ss, nil } @@ -878,9 +1501,6 @@ func (ss *ServerSession) initialized(ctx context.Context, params *InitializedPar ss.server.opts.Logger.Error("duplicate initialized notification") return nil, fmt.Errorf("duplicate %q received", notificationInitialized) } - if ss.server.opts.KeepAlive > 0 { - ss.startKeepalive(ss.server.opts.KeepAlive) - } if h := ss.server.opts.InitializedHandler; h != nil { h(ctx, serverRequestFor(ss, params)) } @@ -910,6 +1530,13 @@ func (ss *ServerSession) NotifyProgress(ctx context.Context, params *ProgressNot return handleNotify(ctx, notificationProgress, newServerRequest(ss, orZero[Params](params))) } +// notifySubscriptionAcked sends a "notifications/subscriptions/acknowledged" +// notification on the listen stream represented by this session, indicating +// the subscription filter the server accepted (SEP-2575). +func (ss *ServerSession) notifySubscriptionAcked(ctx context.Context, params *SubscriptionsAcknowledgedParams) error { + return handleNotify(ctx, notificationSubscriptionsAck, newServerRequest(ss, orZero[Params](params))) +} + func newServerRequest[P Params](ss *ServerSession, params P) *ServerRequest[P] { return &ServerRequest[P]{Session: ss, Params: params} } @@ -930,10 +1557,22 @@ type ServerSession struct { server *Server conn *jsonrpc2.Connection mcpConn Connection - keepaliveCancel context.CancelFunc // TODO: theory around why keepaliveCancel need not be guarded + keepaliveCancel context.CancelFunc + + // supportedVersions is the subset of [Server.protocolVersions] that the + // transport can actually serve, computed once at connection time from + // [ProtocolVersionSupporter] (if implemented by the transport) and used by + // the SEP-2575 server/discover handler. + supportedVersions []string mu sync.Mutex state ServerSessionState + // listenIDs holds the request IDs of in-flight subscriptions/listen + // handlers on this session. subscriptions/listen parks on ctx.Done until + // cancelled, so ServerSession.Close must cancel these contexts explicitly + // via jsonrpc2.Connection.Cancel to avoid deadlocking on the jsonrpc2 + // drain. See modelcontextprotocol/go-sdk#1160. + listenIDs []jsonrpc.ID } func (ss *ServerSession) updateState(mut func(*ServerSessionState)) { @@ -980,6 +1619,23 @@ func (ss *ServerSession) ID() string { return "" } +// assertServerInitiatedRequestAllowed returns an error when the session is +// negotiated at protocol version >= 2026-07-28, where the spec (SEP-2322 / +// SEP-2575) forbids server-initiated JSON-RPC requests for elicitation, +// sampling, and roots: those interactions MUST be embedded as [InputRequests] +// in an [InputRequiredResult] returned from a handler for one of the multi +// round-trip methods (`tools/call`, `prompts/get`, `resources/read`). +func (ss *ServerSession) assertServerInitiatedRequestAllowed(method string) error { + if iparams := ss.InitializeParams(); iparams != nil && + iparams.ProtocolVersion >= protocolVersion20260728 { + return fmt.Errorf( + "%q cannot be sent while serving a request on protocol version %s: "+ + "return an InputRequests map instead (multi round-trip requests, SEP-2322)", + method, iparams.ProtocolVersion) + } + return nil +} + // Ping pings the client. func (ss *ServerSession) Ping(ctx context.Context, params *PingParams) error { _, err := handleSend[*emptyResult](ctx, methodPing, newServerRequest(ss, orZero[Params](params))) @@ -987,18 +1643,40 @@ func (ss *ServerSession) Ping(ctx context.Context, params *PingParams) error { } // ListRoots lists the client roots. +// +// Deprecated: the roots feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). Migrate to passing paths via tool parameters, +// resource URIs, or configuration. See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) { if err := ss.checkInitialized(methodListRoots); err != nil { return nil, err } + if err := ss.assertServerInitiatedRequestAllowed(methodListRoots); err != nil { + return nil, err + } return handleSend[*ListRootsResult](ctx, methodListRoots, newServerRequest(ss, orZero[Params](params))) } // CreateMessage sends a sampling request to the client. +// +// If the client returns multiple content blocks (e.g. parallel tool calls), +// CreateMessage returns an error. Use [ServerSession.CreateMessageWithTools] +// for tool-enabled sampling. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). Migrate to calling LLM provider APIs directly +// from your server. See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessageParams) (*CreateMessageResult, error) { if err := ss.checkInitialized(methodCreateMessage); err != nil { return nil, err } + if err := ss.assertServerInitiatedRequestAllowed(methodCreateMessage); err != nil { + return nil, err + } if params == nil { params = &CreateMessageParams{Messages: []*SamplingMessage{}} } @@ -1007,7 +1685,53 @@ func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessag p2.Messages = []*SamplingMessage{} // avoid JSON "null" params = &p2 } - return handleSend[*CreateMessageResult](ctx, methodCreateMessage, newServerRequest(ss, orZero[Params](params))) + res, err := handleSend[*CreateMessageWithToolsResult](ctx, methodCreateMessage, newServerRequest(ss, orZero[Params](params))) + if err != nil { + return nil, err + } + // Downconvert to singular content. + if len(res.Content) > 1 { + return nil, fmt.Errorf("CreateMessage result has %d content blocks; use CreateMessageWithTools for multiple content", len(res.Content)) + } + var content Content + if len(res.Content) > 0 { + content = res.Content[0] + } + return &CreateMessageResult{ + Meta: res.Meta, + Content: content, + Model: res.Model, + Role: res.Role, + StopReason: res.StopReason, + }, nil +} + +// CreateMessageWithTools sends a sampling request with tools to the client, +// returning a [CreateMessageWithToolsResult] that supports array content +// (for parallel tool calls). Use this instead of [ServerSession.CreateMessage] +// when the request includes tools. +// +// Deprecated: the sampling feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). Migrate to calling LLM provider APIs directly +// from your server. See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. +func (ss *ServerSession) CreateMessageWithTools(ctx context.Context, params *CreateMessageWithToolsParams) (*CreateMessageWithToolsResult, error) { + if err := ss.checkInitialized(methodCreateMessage); err != nil { + return nil, err + } + if err := ss.assertServerInitiatedRequestAllowed(methodCreateMessage); err != nil { + return nil, err + } + if params == nil { + params = &CreateMessageWithToolsParams{Messages: []*SamplingMessageV2{}} + } + if params.Messages == nil { + p2 := *params + p2.Messages = []*SamplingMessageV2{} // avoid JSON "null" + params = &p2 + } + return handleSend[*CreateMessageWithToolsResult](ctx, methodCreateMessage, newServerRequest(ss, orZero[Params](params))) } // Elicit sends an elicitation request to the client asking for user input. @@ -1015,16 +1739,97 @@ func (ss *ServerSession) Elicit(ctx context.Context, params *ElicitParams) (*Eli if err := ss.checkInitialized(methodElicit); err != nil { return nil, err } - return handleSend[*ElicitResult](ctx, methodElicit, newServerRequest(ss, orZero[Params](params))) + if err := ss.assertServerInitiatedRequestAllowed(methodElicit); err != nil { + return nil, err + } + if params == nil { + return nil, fmt.Errorf("%w: params cannot be nil", jsonrpc2.ErrInvalidParams) + } + + params = params.inferElicitMode() + + if iparams := ss.InitializeParams(); iparams == nil || iparams.Capabilities == nil || iparams.Capabilities.Elicitation == nil { + return nil, fmt.Errorf("client does not support elicitation") + } + caps := ss.InitializeParams().Capabilities.Elicitation + switch params.Mode { + case "form": + if caps.Form == nil && caps.URL != nil { + // Note: if both 'Form' and 'URL' are nil, we assume the client supports + // form elicitation for backward compatibility. + return nil, errors.New(`client does not support "form" elicitation`) + } + case "url": + if caps.URL == nil { + return nil, errors.New(`client does not support "url" elicitation`) + } + } + + res, err := handleSend[*ElicitResult](ctx, methodElicit, newServerRequest(ss, orZero[Params](params))) + if err != nil { + return nil, err + } + + if res.Action != "accept" || res.Content == nil { + return res, nil + } + + if params.RequestedSchema == nil { + return res, nil + } + schema, err := validateElicitSchema(params.RequestedSchema) + if err != nil { + return nil, err + } + if schema == nil { + return res, nil + } + + resolved, err := schema.Resolve(nil) + if err != nil { + return nil, err + } + if err := resolved.Validate(res.Content); err != nil { + return nil, fmt.Errorf("elicitation result content does not match requested schema: %v", err) + } + err = resolved.ApplyDefaults(&res.Content) + if err != nil { + return nil, fmt.Errorf("failed to apply schema defaults to elicitation result: %v", err) + } + + return res, nil +} + +// NotifyElicitationComplete tells the client that the out-of-band ("url" mode) +// elicitation identified by params.ElicitationID has finished. +func (ss *ServerSession) NotifyElicitationComplete(ctx context.Context, params *ElicitationCompleteParams) error { + if params == nil { + return fmt.Errorf("%w: params cannot be nil", jsonrpc2.ErrInvalidParams) + } + if params.ElicitationID == "" { + return fmt.Errorf("%w: ElicitationID cannot be empty", jsonrpc2.ErrInvalidParams) + } + return handleNotify(ctx, notificationElicitationComplete, newServerRequest(ss, orZero[Params](params))) } // Log sends a log message to the client. -// The message is not sent if the client has not called SetLevel, or if its level -// is below that of the last SetLevel. +// +// For new-protocol (>= 2026-07-28) requests, the level is taken from the +// originating request's `_meta` field (SEP-2575); an absent or empty value +// suppresses the message per spec. For old-protocol requests, the level is +// taken from the session state set via `logging/setLevel`. +// +// Deprecated: the logging feature is deprecated as of protocol version +// 2026-07-28 (SEP-2577). It remains functional during the deprecation window +// (at least twelve months). See +// https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. func (ss *ServerSession) Log(ctx context.Context, params *LoggingMessageParams) error { - ss.mu.Lock() - logLevel := ss.state.LogLevel - ss.mu.Unlock() + logLevel, ok := logLevelFromContext(ctx) + if !ok { + ss.mu.Lock() + logLevel = ss.state.LogLevel + ss.mu.Unlock() + } if logLevel == "" { // The spec is unclear, but seems to imply that no log messages are sent until the client // sets the level. @@ -1061,6 +1866,9 @@ func (s *Server) AddSendingMiddleware(middleware ...Middleware) { // // Receiving middleware is called when a request is received. It is useful for tasks // such as authentication, request logging and metrics. +// +// A received message need not carry params: use [HasParams] before inspecting +// [Request.GetParams]. func (s *Server) AddReceivingMiddleware(middleware ...Middleware) { s.mu.Lock() defer s.mu.Unlock() @@ -1074,7 +1882,8 @@ func (s *Server) AddReceivingMiddleware(middleware ...Middleware) { // curating these method flags. var serverMethodInfos = map[string]methodInfo{ methodComplete: newServerMethodInfo(serverMethod((*Server).complete), 0), - methodInitialize: newServerMethodInfo(serverSessionMethod((*ServerSession).initialize), 0), + methodDiscover: newServerMethodInfo(serverMethod((*Server).discover), missingParamsOK), + methodInitialize: initializeMethodInfo(), methodPing: newServerMethodInfo(serverSessionMethod((*ServerSession).ping), missingParamsOK), methodListPrompts: newServerMethodInfo(serverMethod((*Server).listPrompts), missingParamsOK), methodGetPrompt: newServerMethodInfo(serverMethod((*Server).getPrompt), 0), @@ -1085,6 +1894,7 @@ var serverMethodInfos = map[string]methodInfo{ methodReadResource: newServerMethodInfo(serverMethod((*Server).readResource), 0), methodSetLevel: newServerMethodInfo(serverSessionMethod((*ServerSession).setLevel), 0), methodSubscribe: newServerMethodInfo(serverMethod((*Server).subscribe), 0), + methodSubscriptionsListen: newServerMethodInfo(serverMethod((*Server).subscriptionsListen), 0), methodUnsubscribe: newServerMethodInfo(serverMethod((*Server).unsubscribe), 0), notificationCancelled: newServerMethodInfo(serverSessionMethod((*ServerSession).cancel), notification|missingParamsOK), notificationInitialized: newServerMethodInfo(serverSessionMethod((*ServerSession).initialized), notification|missingParamsOK), @@ -1092,9 +1902,36 @@ var serverMethodInfos = map[string]methodInfo{ notificationProgress: newServerMethodInfo(serverSessionMethod((*ServerSession).callProgressNotificationHandler), notification), } +// initializeMethodInfo handles the workaround for #607: we must set +// params.Capabilities.RootsV2. +func initializeMethodInfo() methodInfo { + info := newServerMethodInfo(serverSessionMethod((*ServerSession).initialize), 0) + info.unmarshalParams = func(m json.RawMessage) (Params, error) { + var params *initializeParamsV2 + if m != nil { + if err := internaljson.Unmarshal(m, ¶ms); err != nil { + return nil, fmt.Errorf("unmarshaling %q into a %T: %w", m, params, err) + } + } + if params == nil { + return nil, fmt.Errorf(`missing required "params"`) + } + return params.toV1(), nil + } + return info +} + func (ss *ServerSession) sendingMethodInfos() map[string]methodInfo { return clientMethodInfos } -func (ss *ServerSession) receivingMethodInfos() map[string]methodInfo { return serverMethodInfos } +func (s *Server) receivingMethodInfos() map[string]methodInfo { + s.mu.Lock() + defer s.mu.Unlock() + return s.receiveMethods +} + +func (ss *ServerSession) receivingMethodInfos() map[string]methodInfo { + return ss.server.receivingMethodInfos() +} func (ss *ServerSession) sendingMethodHandler() MethodHandler { s := ss.server @@ -1119,18 +1956,65 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any, initialized := ss.state.InitializeParams != nil ss.mu.Unlock() - // From the spec: - // "The client SHOULD NOT send requests other than pings before the server - // has responded to the initialize request." - switch req.Method { - case methodInitialize, methodPing, notificationInitialized: - default: - if !initialized { - ss.server.opts.Logger.Error("method invalid during initialization", "method", req.Method) - return nil, fmt.Errorf("method %q is invalid during session initialization", req.Method) + // Per-request protocol detection (SEP-2575): if the request carries + // `io.modelcontextprotocol/protocolVersion` in its `_meta` field, it + // follows the new sessionless protocol. The initialization gate is + // skipped for such requests. + validatedMeta, perRequestErr := validateRequestMeta(req) + if perRequestErr != nil { + return nil, perRequestErr + } + + if validatedMeta.usesNewProtocol && + !slices.Contains(ss.server.protocolVersions, validatedMeta.initializeParams.ProtocolVersion) { + data, _ := json.Marshal(UnsupportedProtocolVersionData{ + Supported: ss.server.protocolVersions, + Requested: validatedMeta.initializeParams.ProtocolVersion, + }) + return nil, &jsonrpc.Error{ + Code: CodeUnsupportedProtocolVersion, + Message: "unsupported protocol version", + Data: data, } } + switch req.Method { + case methodInitialize, methodPing, notificationInitialized, notificationRootsListChanged, methodSetLevel, methodSubscribe, methodUnsubscribe: + if validatedMeta.usesNewProtocol { + ss.server.opts.Logger.Error("method removed in the new protocol", "method", req.Method) + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeMethodNotFound, + Message: fmt.Sprintf("%q is not supported in the new protocol", req.Method), + } + } + case methodDiscover: + // In case of methodDiscover call the state.initializeParams is populated + // within the discover handle function to make sure the method is supported + // when the user is probing a pre-2026-07-28 server. + if !validatedMeta.usesNewProtocol { + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeMethodNotFound, + Message: fmt.Sprintf("%q is only supported in protocol version >= %s", req.Method, protocolVersion20260728), + } + } + default: + if !initialized && validatedMeta.usesNewProtocol && validatedMeta.initializeParams != nil { + ss.updateState(func(state *ServerSessionState) { + state.InitializeParams = validatedMeta.initializeParams + }) + } + } + + // In legacy protocol versions, a client cannot send requests other than + // pings before the server has responded to the initialize request. A + // new-protocol request is exempt: a SEP-2575 session has no 'initialize' + // to wait for, and the request itself says which protocol it speaks. + if !initialized && !validatedMeta.usesNewProtocol && req.IsCall() && + req.Method != methodInitialize && req.Method != methodPing { + ss.server.opts.Logger.Error("method invalid during initialization", "method", req.Method) + return nil, fmt.Errorf("method %q is invalid during session initialization", req.Method) + } + // modelcontextprotocol/go-sdk#26: handle calls asynchronously, and // notifications synchronously, except for 'initialize' which shouldn't be // asynchronous to other @@ -1142,24 +2026,101 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any, // server->client calls and notifications to the incoming request from which // they originated. See [idContextKey] for details. ctx = context.WithValue(ctx, idContextKey{}, req.ID) - return handleReceive(ctx, ss, req) + if validatedMeta.usesNewProtocol { + ctx = context.WithValue(ctx, logLevelContextKey{}, validatedMeta.logLevel) + } + + // subscriptions/listen parks on ctx.Done until the peer cancels it (or the + // underlying reader breaks). Track the request ID so ServerSession.Close + // can cancel the in-flight handler via jsonrpc2.Connection.Cancel and + // avoid deadlocking on the jsonrpc2 drain. + if req.Method == methodSubscriptionsListen { + ss.mu.Lock() + ss.listenIDs = append(ss.listenIDs, req.ID) + ss.mu.Unlock() + // The listen completes when the handler returns (peer cancellation, + // stream break, or error); drop the ID so completed listens don't + // accumulate in the slice indefinitely. + defer func() { + ss.mu.Lock() + for i, id := range ss.listenIDs { + if id == req.ID { + ss.listenIDs = append(ss.listenIDs[:i], ss.listenIDs[i+1:]...) + break + } + } + ss.mu.Unlock() + }() + } + + res, err := handleReceive(ctx, ss, req) + if err != nil { + return nil, err + } + // A middleware can return a typed nil value that satisfies the Result + // interface. Annotating that value panics. + if validatedMeta.usesNewProtocol && res != nil && !res.isNil() { + annotateResultType(res) + annotateServerInfo(res, ss.server.impl) + } + return res, nil } -func (ss *ServerSession) InitializeParams() *InitializeParams { return ss.state.InitializeParams } +// annotateServerInfo sets [MetaKeyServerInfo] on the result's `_meta` unless +// it is already present. Per SEP-2575, servers should identify themselves in +// each result's `_meta`. +func annotateServerInfo(res Result, impl *Implementation) { + if res == nil || impl == nil { + return + } + if _, isEmpty := res.(*emptyResult); isEmpty { + return + } + m := res.GetMeta() + if m == nil { + m = map[string]any{} + } + if _, ok := m[MetaKeyServerInfo]; ok { + return + } + m[MetaKeyServerInfo] = impl + res.SetMeta(m) +} + +// InitializeParams returns the InitializeParams provided during the client's +// initial connection. +func (ss *ServerSession) InitializeParams() *InitializeParams { + ss.mu.Lock() + defer ss.mu.Unlock() + return ss.state.InitializeParams +} func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParams) (*InitializeResult, error) { if params == nil { return nil, fmt.Errorf("%w: \"params\" must be be provided", jsonrpc2.ErrInvalidParams) } + var ( + wasInit bool + negotiated string + ) ss.updateState(func(state *ServerSessionState) { - state.InitializeParams = params + wasInit = state.InitializeParams != nil + if !wasInit { + state.InitializeParams = params + state.NegotiatedProtocolVersion = negotiatedVersion(params.ProtocolVersion, ss.server.protocolVersions) + } + negotiated = state.NegotiatedProtocolVersion }) + if wasInit { + ss.server.opts.Logger.Error("duplicate initialize request") + return nil, fmt.Errorf("duplicate %q received", methodInitialize) + } s := ss.server return &InitializeResult{ // TODO(rfindley): alter behavior when falling back to an older version: // reject unsupported features. - ProtocolVersion: negotiatedVersion(params.ProtocolVersion), + ProtocolVersion: negotiated, Capabilities: s.capabilities(), Instructions: s.opts.Instructions, ServerInfo: s.impl, @@ -1195,12 +2156,25 @@ func (ss *ServerSession) setLevel(_ context.Context, params *SetLoggingLevelPara func (ss *ServerSession) Close() error { if ss.keepaliveCancel != nil { // Note: keepaliveCancel access is safe without a mutex because: - // 1. keepaliveCancel is only written once during startKeepalive (happens-before all Close calls) + // 1. keepaliveCancel is only written once during Server.Connect (through startKeepalive), + // which happens before any code that may call Close from another goroutine // 2. context.CancelFunc is safe to call multiple times and from multiple goroutines // 3. The keepalive goroutine calls Close on ping failure, but this is safe since // Close is idempotent and conn.Close() handles concurrent calls correctly ss.keepaliveCancel() } + + // Unblock any in-flight subscriptions/listen handlers, which otherwise park + // on ctx.Done and would deadlock conn.Close (which waits for in-flight + // requests to drain). + ss.mu.Lock() + ids := ss.listenIDs + ss.listenIDs = nil + ss.mu.Unlock() + for _, id := range ids { + ss.conn.Cancel(id) + } + err := ss.conn.Close() if ss.onClose != nil && ss.calledOnClose.CompareAndSwap(false, true) { @@ -1217,7 +2191,7 @@ func (ss *ServerSession) Wait() error { // startKeepalive starts the keepalive mechanism for this server session. func (ss *ServerSession) startKeepalive(interval time.Duration) { - startKeepalive(ss, interval, &ss.keepaliveCancel) + startKeepalive(ss, interval, ss.server.opts.KeepAliveFailureThreshold, &ss.keepaliveCancel, ss.server.opts.Logger) } // pageToken is the internal structure for the opaque pagination cursor. @@ -1295,3 +2269,55 @@ func paginateList[P listParams, R listResult[T], T any](fs *featureSet[T], pageS *res.nextCursorPtr() = nextCursor return res, nil } + +// AddReceivingCustomMethod registers a handler for a custom (non-standard) +// JSON-RPC method on the server. +// +// When a client sends a request with the given method name, the params will be +// unmarshaled into P, the handler will be called, and the returned R will be +// marshaled as the JSON-RPC result. +// +// Custom methods go through the server's middleware chain just like standard +// MCP methods (tools/call, prompts/list, etc.). +// +// P and R must implement [Params] and [Result] respectively, which is most +// easily done by embedding [ParamsBase] and [ResultBase]: +// +// type SearchParams struct { +// mcp.ParamsBase +// Query string `json:"query"` +// } +// +// type SearchResult struct { +// mcp.ResultBase +// Hits []string `json:"hits"` +// } +// +// if err := mcp.AddReceivingCustomMethod(server, "acme/search", +// func(ctx context.Context, ss *mcp.ServerSession, params *SearchParams) (*SearchResult, error) { +// return &SearchResult{Hits: []string{"result"}}, nil +// }); err != nil { +// return err +// } +// +// AddReceivingCustomMethod returns an error if method is the name of a +// standard MCP method. Registering the same custom method twice replaces the +// previous handler. +func AddReceivingCustomMethod[P paramsPtr[T], R Result, T any]( + s *Server, + method string, + handler func(ctx context.Context, ss *ServerSession, params P) (R, error), +) error { + if _, ok := serverMethodInfos[method]; ok { + return fmt.Errorf("mcp: AddReceivingCustomMethod: %q shadows a standard MCP method", method) + } + + typed := typedServerMethodHandler[P, R](func(ctx context.Context, req *ServerRequest[P]) (R, error) { + return handler(ctx, req.Session, req.Params) + }) + + s.mu.Lock() + defer s.mu.Unlock() + s.receiveMethods[method] = newServerMethodInfo(typed, missingParamsOK) + return nil +} diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index db04920..ae4318f 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -2,6 +2,8 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. +//lint:file-ignore SA1019 examples exercise deprecated SEP-2577 logging APIs for demonstration. + package mcp_test import ( @@ -11,6 +13,7 @@ import ( "log/slog" "sync/atomic" + "github.com/google/jsonschema-go/jsonschema" "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -85,6 +88,77 @@ func Example_prompts() { // !-prompts +// !+promptcontent + +func Example_promptContent() { + ctx := context.Background() + + const styleGuide = "- Prefer clarity over cleverness.\n" + screenshotPNG := []byte("\x89PNG\r\n\x1a\n") + + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + + s.AddResource(&mcp.Resource{URI: "doc://style-guide", MIMEType: "text/markdown"}, + func(context.Context, *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{ + Contents: []*mcp.ResourceContents{{ + URI: "doc://style-guide", + MIMEType: "text/markdown", + Text: styleGuide, + }}, + }, nil + }) + + s.AddPrompt(&mcp.Prompt{Name: "review_screenshot"}, + func(context.Context, *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{ + Messages: []*mcp.PromptMessage{ + {Role: "user", Content: &mcp.EmbeddedResource{ + Resource: &mcp.ResourceContents{ + URI: "doc://style-guide", + MIMEType: "text/markdown", + Text: styleGuide, + }, + }}, + {Role: "user", Content: &mcp.ImageContent{Data: screenshotPNG, MIMEType: "image/png"}}, + {Role: "user", Content: &mcp.TextContent{Text: "Review the screenshot against the style guide."}}, + }, + }, nil + }) + + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) + t1, t2 := mcp.NewInMemoryTransports() + if _, err := s.Connect(ctx, t1, nil); err != nil { + log.Fatal(err) + } + cs, err := c.Connect(ctx, t2, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + res, err := cs.GetPrompt(ctx, &mcp.GetPromptParams{Name: "review_screenshot"}) + if err != nil { + log.Fatal(err) + } + for _, msg := range res.Messages { + switch content := msg.Content.(type) { + case *mcp.TextContent: + fmt.Println(msg.Role, "text:", content.Text) + case *mcp.ImageContent: + fmt.Printf("%s image: %s, %d bytes\n", msg.Role, content.MIMEType, len(content.Data)) + case *mcp.EmbeddedResource: + fmt.Printf("%s embedded resource: %s (%s)\n", msg.Role, content.Resource.URI, content.Resource.MIMEType) + } + } + // Output: + // user embedded resource: doc://style-guide (text/markdown) + // user image: image/png, 8 bytes + // user text: Review the screenshot against the style guide. +} + +// !-promptcontent + // !+logging func Example_logging() { @@ -218,3 +292,269 @@ func Example_resources() { } // !-resources + +// !+listchanged + +func Example_listChanged() { + ctx := context.Background() + + changed := make(chan string, 3) + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, &mcp.ClientOptions{ + ToolListChangedHandler: func(context.Context, *mcp.ToolListChangedRequest) { + changed <- "tools" + }, + PromptListChangedHandler: func(context.Context, *mcp.PromptListChangedRequest) { + changed <- "prompts" + }, + ResourceListChangedHandler: func(context.Context, *mcp.ResourceListChangedRequest) { + changed <- "resources" + }, + }) + + // Nothing is registered before Connect, so the capabilities that carry these + // notifications have to be declared explicitly. + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, &mcp.ServerOptions{ + HasTools: true, + HasPrompts: true, + HasResources: true, + }) + + t1, t2 := mcp.NewInMemoryTransports() + if _, err := s.Connect(ctx, t1, nil); err != nil { + log.Fatal(err) + } + cs, err := c.Connect(ctx, t2, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + mcp.AddTool(s, &mcp.Tool{Name: "greet"}, func(context.Context, *mcp.CallToolRequest, struct{}) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{}, nil, nil + }) + fmt.Println("changed:", <-changed) + + s.AddPrompt(&mcp.Prompt{Name: "review"}, func(context.Context, *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{}, nil + }) + fmt.Println("changed:", <-changed) + + s.AddResource(&mcp.Resource{URI: "file:///a"}, func(context.Context, *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{}, nil + }) + fmt.Println("changed:", <-changed) + + // Removing a feature notifies the client just as adding one does. + s.RemoveTools("greet") + fmt.Println("changed:", <-changed) + + // Output: + // changed: tools + // changed: prompts + // changed: resources + // changed: tools +} + +// !-listchanged + +// !+subscribe + +func Example_resourceSubscription() { + ctx := context.Background() + + config := "theme=light\n" + + updated := make(chan string, 1) + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, &mcp.ClientOptions{ + // The notification carries only the URI, so read the resource again to + // see what it now holds. + ResourceUpdatedHandler: func(context.Context, *mcp.ResourceUpdatedNotificationRequest) { + updated <- "config://app" + }, + }) + + subscribed := make(chan string, 2) + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, &mcp.ServerOptions{ + SubscribeHandler: func(_ context.Context, req *mcp.SubscribeRequest) error { + subscribed <- "subscribed to " + req.Params.URI + return nil + }, + UnsubscribeHandler: func(_ context.Context, req *mcp.UnsubscribeRequest) error { + subscribed <- "unsubscribed from " + req.Params.URI + return nil + }, + }) + s.AddResource(&mcp.Resource{URI: "config://app"}, + func(context.Context, *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{ + Contents: []*mcp.ResourceContents{{URI: "config://app", Text: config}}, + }, nil + }) + + t1, t2 := mcp.NewInMemoryTransports() + if _, err := s.Connect(ctx, t1, nil); err != nil { + log.Fatal(err) + } + cs, err := c.Connect(ctx, t2, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + if err := cs.Subscribe(ctx, &mcp.SubscribeParams{URI: "config://app"}); err != nil { + log.Fatal(err) + } + fmt.Println(<-subscribed) + + config = "theme=dark\n" + if err := s.ResourceUpdated(ctx, &mcp.ResourceUpdatedNotificationParams{URI: "config://app"}); err != nil { + log.Fatal(err) + } + uri := <-updated + + res, err := cs.ReadResource(ctx, &mcp.ReadResourceParams{URI: uri}) + if err != nil { + log.Fatal(err) + } + fmt.Print("re-read: ", res.Contents[0].Text) + + if err := cs.Unsubscribe(ctx, &mcp.UnsubscribeParams{URI: "config://app"}); err != nil { + log.Fatal(err) + } + fmt.Println(<-subscribed) + + // Output: + // subscribed to config://app + // re-read: theme=dark + // unsubscribed from config://app +} + +// !-subscribe + +// !+binaryresource + +func Example_binaryResource() { + ctx := context.Background() + + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + s.AddResource(&mcp.Resource{URI: "file:///logo.png", MIMEType: "image/png"}, + func(context.Context, *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{ + Contents: []*mcp.ResourceContents{{ + URI: "file:///logo.png", + MIMEType: "image/png", + Blob: []byte("\x89PNG\r\n\x1a\n"), + }}, + }, nil + }) + + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) + t1, t2 := mcp.NewInMemoryTransports() + if _, err := s.Connect(ctx, t1, nil); err != nil { + log.Fatal(err) + } + cs, err := c.Connect(ctx, t2, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + res, err := cs.ReadResource(ctx, &mcp.ReadResourceParams{URI: "file:///logo.png"}) + if err != nil { + log.Fatal(err) + } + contents := res.Contents[0] + fmt.Printf("%s: %d bytes of %s, text is %q\n", + contents.URI, len(contents.Blob), contents.MIMEType, contents.Text) + + // Output: + // file:///logo.png: 8 bytes of image/png, text is "" +} + +// !-binaryresource + +// !+mrtr + +// Example_mrtr demonstrates the [Multi Round-Trip Requests] pattern +// (SEP-2322). A tool handler signals "I need more information from the user" +// by returning a [mcp.CallToolResult] whose `InputRequests` field carries the +// requests for the additional information. Each request can be an +// [mcp.ElicitParams] (ask the user a question), +// [mcp.CreateMessageParams] (sample from the client's LLM), or +// [mcp.ListRootsParams] (list the client's roots). +// +// On protocol version `2026-07-28` and later, the client's +// `clientMultiRoundTripMiddleware` fulfils each input request from the +// configured handler and retries the original call transparently. On earlier +// protocol versions, the server's `serverMultiRoundTripMiddleware` performs +// the equivalent dance from the server side by calling the legacy +// server-to-client API (`ServerSession.Elicit`, `CreateMessage`, or +// `ListRoots`) and re-invoking the handler with the response. Either way, +// the user-facing client call site sees only the final result. +// +// [Multi Round-Trip Requests]: https://modelcontextprotocol.io/specification/draft/basic/patterns#multi-round-trip-requests +func Example_mrtr() { + ctx := context.Background() + + // Server: a "greet" tool that asks the user for their name before + // returning a greeting. The handler runs twice per logical call: + // once to issue the elicitation, once to consume the response. + s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(s, &mcp.Tool{Name: "greet"}, func(_ context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + if len(req.Params.InputResponses) == 0 { + // First call: ask the user for their name. The map key + // ("user_name") is a label the handler chooses; the client must + // echo it back in InputResponses. The wire-level method name + // ("elicitation/create") is derived from the value's Go type. + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{ + "user_name": &mcp.ElicitParams{ + Message: "What's your name?", + RequestedSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "name": {Type: "string"}, + }, + }, + }, + }, + // RequestState is an opaque token the client echoes back on + // the retry, letting the handler resume its work without + // per-session storage. + RequestState: "step=1", + }, nil, nil + } + // Retry: read the elicitation response and produce the final greeting. + name := req.Params.InputResponses["user_name"].(*mcp.ElicitResult).Content["name"].(string) + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "Hello " + name}}, + }, nil, nil + }) + + // Client: declares an elicitation handler. The SDK's MRTR middleware + // uses it to fulfil any input request the tool handler returns. + c := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, &mcp.ClientOptions{ + ElicitationHandler: func(_ context.Context, _ *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + return &mcp.ElicitResult{Action: "accept", Content: map[string]any{"name": "MCP Go"}}, nil + }, + }) + + ct, st := mcp.NewInMemoryTransports() + if _, err := s.Connect(ctx, st, nil); err != nil { + log.Fatal(err) + } + cs, err := c.Connect(ctx, ct, nil) + if err != nil { + log.Fatal(err) + } + defer cs.Close() + + // Single call site. The MRTR round trip is invisible to the caller. + res, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: "greet"}) + if err != nil { + log.Fatal(err) + } + fmt.Println(res.Content[0].(*mcp.TextContent).Text) +} + +// !-mrtr diff --git a/mcp/server_test.go b/mcp/server_test.go index ec6114d..4d09cbc 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -5,9 +5,13 @@ package mcp import ( + "bytes" "context" "encoding/json" + "errors" + "fmt" "log" + "log/slog" "slices" "strings" "testing" @@ -15,6 +19,8 @@ import ( "github.com/google/go-cmp/cmp" "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) type testItem struct { @@ -242,14 +248,14 @@ func TestServerCapabilities(t *testing.T) { wantCapabilities *ServerCapabilities }{ { - name: "No capabilities", + name: "no capabilities", configureServer: func(s *Server) {}, wantCapabilities: &ServerCapabilities{ Logging: &LoggingCapabilities{}, }, }, { - name: "With prompts", + name: "with prompts", configureServer: func(s *Server) { s.AddPrompt(&Prompt{Name: "p"}, nil) }, @@ -259,7 +265,7 @@ func TestServerCapabilities(t *testing.T) { }, }, { - name: "With resources", + name: "with resources", configureServer: func(s *Server) { s.AddResource(&Resource{URI: "file:///r"}, nil) }, @@ -269,7 +275,7 @@ func TestServerCapabilities(t *testing.T) { }, }, { - name: "With resource templates", + name: "with resource templates", configureServer: func(s *Server) { s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil) }, @@ -279,7 +285,7 @@ func TestServerCapabilities(t *testing.T) { }, }, { - name: "With resource subscriptions", + name: "with resource subscriptions", configureServer: func(s *Server) { s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil) }, @@ -297,7 +303,7 @@ func TestServerCapabilities(t *testing.T) { }, }, { - name: "With tools", + name: "with tools", configureServer: func(s *Server) { s.AddTool(tool, nil) }, @@ -307,7 +313,7 @@ func TestServerCapabilities(t *testing.T) { }, }, { - name: "With completions", + name: "with completions", configureServer: func(s *Server) {}, serverOpts: ServerOptions{ CompletionHandler: func(context.Context, *CompleteRequest) (*CompleteResult, error) { @@ -320,7 +326,7 @@ func TestServerCapabilities(t *testing.T) { }, }, { - name: "With all capabilities", + name: "all capabilities", configureServer: func(s *Server) { s.AddPrompt(&Prompt{Name: "p"}, nil) s.AddResource(&Resource{URI: "file:///r"}, nil) @@ -347,7 +353,7 @@ func TestServerCapabilities(t *testing.T) { }, }, { - name: "With initial capabilities", + name: "has features", configureServer: func(s *Server) {}, serverOpts: ServerOptions{ HasPrompts: true, @@ -361,6 +367,102 @@ func TestServerCapabilities(t *testing.T) { Tools: &ToolCapabilities{ListChanged: true}, }, }, + { + name: "empty capabilities", + configureServer: func(s *Server) {}, + serverOpts: ServerOptions{ + Capabilities: &ServerCapabilities{}, + }, + wantCapabilities: &ServerCapabilities{}, + }, + { + name: "no logging", + configureServer: func(s *Server) {}, + serverOpts: ServerOptions{ + Capabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: true}, + }, + }, + wantCapabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: true}, + }, + }, + { + name: "no list", + configureServer: func(s *Server) {}, + serverOpts: ServerOptions{ + Capabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: false}, + Prompts: &PromptCapabilities{ListChanged: false}, + }, + }, + wantCapabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: false}, + Prompts: &PromptCapabilities{ListChanged: false}, + }, + }, + { + name: "adding tools-list", + configureServer: func(s *Server) { + s.AddTool(tool, nil) + }, + serverOpts: ServerOptions{ + Capabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + }, + }, + wantCapabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Tools: &ToolCapabilities{ListChanged: true}, + }, + }, + { + name: "adding tools-no list", + configureServer: func(s *Server) { + s.AddTool(tool, nil) + }, + serverOpts: ServerOptions{ + Capabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: false}, + }, + }, + wantCapabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: false}, + }, + }, + { + name: "experimental preserved", + configureServer: func(s *Server) {}, + serverOpts: ServerOptions{ + Capabilities: &ServerCapabilities{ + Experimental: map[string]any{"custom": "value"}, + Logging: &LoggingCapabilities{}, + }, + }, + wantCapabilities: &ServerCapabilities{ + Experimental: map[string]any{"custom": "value"}, + Logging: &LoggingCapabilities{}, + }, + }, + { + name: "extensions preserved", + configureServer: func(s *Server) {}, + serverOpts: func() ServerOptions { + caps := &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + } + caps.AddExtension("io.example/ext1", map[string]any{"key": "value"}) + caps.AddExtension("io.example/ext2", nil) + return ServerOptions{Capabilities: caps} + }(), + wantCapabilities: &ServerCapabilities{ + Extensions: map[string]any{ + "io.example/ext1": map[string]any{"key": "value"}, + "io.example/ext2": map[string]any{}, + }, + Logging: &LoggingCapabilities{}, + }, + }, } for _, tc := range testCases { @@ -409,54 +511,6 @@ func TestServerAddResourceTemplate(t *testing.T) { } } -// TestServerSessionkeepaliveCancelOverwritten is to verify that `ServerSession.keepaliveCancel` is assigned exactly once, -// ensuring that only a single goroutine is responsible for the session's keepalive ping mechanism. -func TestServerSessionkeepaliveCancelOverwritten(t *testing.T) { - // Set KeepAlive to a long duration to ensure the keepalive - // goroutine stays alive for the duration of the test without actually sending - // ping requests, since we don't have a real client connection established. - server := NewServer(testImpl, &ServerOptions{KeepAlive: 5 * time.Second}) - ss := &ServerSession{server: server} - - // 1. Initialize the session. - _, err := ss.initialize(context.Background(), &InitializeParams{}) - if err != nil { - t.Fatalf("ServerSession initialize failed: %v", err) - } - - // 2. Call 'initialized' for the first time. This should start the keepalive mechanism. - _, err = ss.initialized(context.Background(), &InitializedParams{}) - if err != nil { - t.Fatalf("First initialized call failed: %v", err) - } - if ss.keepaliveCancel == nil { - t.Fatalf("expected ServerSession.keepaliveCancel to be set after the first call of initialized") - } - - // Save the cancel function and use defer to ensure resources are cleaned up. - firstCancel := ss.keepaliveCancel - defer firstCancel() - - // 3. Manually set the field to nil. - // Do this to facilitate the test's core assertion. The goal is to verify that - // 'ss.keepaliveCancel' is not assigned a second time. By setting it to nil, - // we can easily check after the next call if a new keepalive goroutine was started. - ss.keepaliveCancel = nil - - // 4. Call 'initialized' for the second time. This should return an error. - _, err = ss.initialized(context.Background(), &InitializedParams{}) - if err == nil { - t.Fatalf("Expected 'duplicate initialized received' error on second call, got nil") - } - - // 5. Re-check the field to ensure it remains nil. - // Since 'initialized' correctly returned an error and did not call - // 'startKeepalive', the field should remain unchanged. - if ss.keepaliveCancel != nil { - t.Fatal("expected ServerSession.keepaliveCancel to be nil after we manually niled it and re-initialized") - } -} - // panicks reports whether f() panics. func panics(f func()) (b bool) { defer func() { @@ -467,7 +521,8 @@ func panics(f func()) (b bool) { } func TestAddTool(t *testing.T) { - // AddTool should panic if In or Out are not JSON objects. + // AddTool should panic if In is not a JSON object. + // Out may be of any type whose inferred JSON Schema is valid. s := NewServer(testImpl, nil) if !panics(func() { AddTool(s, &Tool{Name: "T1"}, func(context.Context, *CallToolRequest, string) (*CallToolResult, any, error) { return nil, nil, nil }) @@ -481,15 +536,359 @@ func TestAddTool(t *testing.T) { }) { t.Error("good In: expected no panic") } - if !panics(func() { - AddTool(s, &Tool{Name: "T2"}, func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, int, error) { + if panics(func() { + AddTool(s, &Tool{Name: "T3"}, func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, int, error) { return nil, 0, nil }) }) { - t.Error("bad Out: expected panic") + t.Error("primitive Out: expected no panic") } } +func TestAddToolNameValidation(t *testing.T) { + tests := []struct { + label string + name string + wantLogContaining string + }{ + { + label: "empty name", + name: "", + wantLogContaining: `tool name cannot be empty`, + }, + { + label: "long name", + name: strings.Repeat("a", 129), + wantLogContaining: "exceeds maximum length of 128 characters", + }, + { + label: "name with spaces", + name: "get user profile", + wantLogContaining: `tool name contains invalid characters: \" \"`, + }, + { + label: "name with multiple invalid chars", + name: "user name@domain,com", + wantLogContaining: `tool name contains invalid characters: \" \", \"@\", \",\"`, + }, + { + label: "name with unicode", + name: "tool-ñame", + wantLogContaining: `tool name contains invalid characters: \"ñ\"`, + }, + { + label: "valid name", + name: "valid-tool_name.123", + wantLogContaining: "", // No log expected + }, + } + for _, test := range tests { + t.Run(test.label, func(t *testing.T) { + var buf bytes.Buffer + s := NewServer(testImpl, &ServerOptions{ + Logger: slog.New(slog.NewTextHandler(&buf, nil)), + }) + + // Use the generic AddTool as it also calls validateToolName. + AddTool(s, &Tool{Name: test.name}, func(context.Context, *CallToolRequest, any) (*CallToolResult, any, error) { + return nil, nil, nil + }) + + logOutput := buf.String() + if test.wantLogContaining != "" { + if !strings.Contains(logOutput, test.wantLogContaining) { + t.Errorf("log output =\n%s\nwant containing %q", logOutput, test.wantLogContaining) + } + } else { + if logOutput != "" { + t.Errorf("expected empty log output, got %q", logOutput) + } + } + }) + } +} + +func TestAddToolNilSchema(t *testing.T) { + var nilSchema *jsonschema.Schema + + panicMsg := func(f func()) (msg string) { + defer func() { + if r := recover(); r != nil { + msg = fmt.Sprintf("%v", r) + } + }() + f() + return msg + } + + // Call s.AddTool directly to exercise the typed-nil checks added to that method. + tests := []struct { + name string + tool *Tool + wantContain string + }{ + { + name: "typed nil InputSchema", + tool: &Tool{Name: "T", InputSchema: nilSchema}, + wantContain: "input schema is nil", + }, + { + name: "typed nil OutputSchema", + tool: &Tool{Name: "T", InputSchema: &jsonschema.Schema{Type: "object"}, OutputSchema: nilSchema}, + wantContain: "output schema is nil", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + s := NewServer(testImpl, nil) + + msg := panicMsg(func() { + s.AddTool(tc.tool, nil) + }) + if msg == "" { + t.Error("expected panic") + } else if !strings.Contains(msg, tc.wantContain) { + t.Errorf("panic message %q does not contain %q", msg, tc.wantContain) + } + }) + } +} + +// TestAddToolNonObjectOutputSchema verifies the low-level +// Server.AddTool accepts an output schema whose root type is not "object" +// (array or primitive) and structured content of the matching shape +// round-trips end-to-end. +func TestAddToolNonObjectOutputSchema(t *testing.T) { + ctx := context.Background() + + for _, tc := range []struct { + name string + outputSchema any + content any + want any + }{ + { + name: "array of objects", + outputSchema: &jsonschema.Schema{Type: "array", Items: &jsonschema.Schema{Type: "object"}}, + content: []any{map[string]any{"id": "u1"}, map[string]any{"id": "u2"}}, + want: []any{map[string]any{"id": "u1"}, map[string]any{"id": "u2"}}, + }, + { + name: "primitive number (map-based schema)", + outputSchema: map[string]any{"type": "number"}, + content: 42.0, + want: json.Number("42"), + }, + { + name: "primitive string (RawMessage schema)", + outputSchema: json.RawMessage(`{"type":"string"}`), + content: "hello", + want: "hello", + }, + } { + t.Run(tc.name, func(t *testing.T) { + server := NewServer(testImpl, nil) + handler := func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{StructuredContent: tc.content}, nil + } + tool := &Tool{ + Name: "t", + InputSchema: &jsonschema.Schema{Type: "object"}, + OutputSchema: tc.outputSchema, + } + // Must not panic. + server.AddTool(tool, handler) + + ct, st := NewInMemoryTransports() + if _, err := server.Connect(ctx, st, nil); err != nil { + t.Fatal(err) + } + client := NewClient(testImpl, nil) + cs, err := client.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + res, err := cs.CallTool(ctx, &CallToolParams{Name: "t"}) + if err != nil { + t.Fatal(err) + } + if res.IsError { + t.Fatalf("unexpected tool error: %v", res.Content) + } + if diff := cmp.Diff(tc.want, res.StructuredContent); diff != "" { + t.Errorf("structured content mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestCallToolStructuredContentPreservesLargeInteger(t *testing.T) { + const want int64 = 9007199254740993 + + server := NewServer(testImpl, nil) + server.AddTool(&Tool{ + Name: "large-integer", + InputSchema: &jsonschema.Schema{Type: "object"}, + }, func(context.Context, *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{StructuredContent: map[string]any{"id": want}}, nil + }) + + clientTransport, serverTransport := NewInMemoryTransports() + serverSession, err := server.Connect(context.Background(), serverTransport, nil) + if err != nil { + t.Fatal(err) + } + defer serverSession.Close() + + client := NewClient(testImpl, nil) + clientSession, err := client.Connect(context.Background(), clientTransport, nil) + if err != nil { + t.Fatal(err) + } + defer clientSession.Close() + + result, err := clientSession.CallTool(context.Background(), &CallToolParams{Name: "large-integer"}) + if err != nil { + t.Fatal(err) + } + structured, ok := result.StructuredContent.(map[string]any) + if !ok { + t.Fatalf("StructuredContent type = %T, want map[string]any", result.StructuredContent) + } + number, ok := structured["id"].(json.Number) + if !ok { + t.Fatalf("id type = %T, want json.Number", structured["id"]) + } + got, err := number.Int64() + if err != nil { + t.Fatal(err) + } + if got != want { + t.Errorf("id = %d, want %d", got, want) + } +} + +// TestAddToolGenericNonObjectOutput verifies SEP-2106 for the generic +// AddTool[In, Out] helper: Out may be a slice or primitive whose inferred +// JSON Schema has a non-object root. +func TestAddToolGenericNonObjectOutput(t *testing.T) { + ctx := context.Background() + + type item struct { + ID string `json:"id"` + } + + t.Run("slice output", func(t *testing.T) { + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "list"}, + func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, []item, error) { + return nil, []item{{ID: "u1"}, {ID: "u2"}}, nil + }) + + ct, st := NewInMemoryTransports() + if _, err := server.Connect(ctx, st, nil); err != nil { + t.Fatal(err) + } + client := NewClient(testImpl, nil) + cs, err := client.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Verify tools/list reports an array-rooted output schema. + // jsonschema-go infers ["null","array"] for slices since nil slices + // are valid; accept either form. + lt, err := cs.ListTools(ctx, nil) + if err != nil { + t.Fatal(err) + } + if len(lt.Tools) != 1 { + t.Fatalf("got %d tools, want 1", len(lt.Tools)) + } + gotSchema, _ := lt.Tools[0].OutputSchema.(map[string]any) + typeOK := false + switch typ := gotSchema["type"].(type) { + case string: + typeOK = typ == "array" + case []any: + for _, e := range typ { + if e == "array" { + typeOK = true + } + } + } + if !typeOK { + t.Errorf("output schema type = %v, want to include %q", gotSchema["type"], "array") + } + + res, err := cs.CallTool(ctx, &CallToolParams{Name: "list"}) + if err != nil { + t.Fatal(err) + } + if res.IsError { + t.Fatalf("unexpected tool error: %v", res.Content) + } + want := []any{map[string]any{"id": "u1"}, map[string]any{"id": "u2"}} + if diff := cmp.Diff(want, res.StructuredContent); diff != "" { + t.Errorf("structured content mismatch (-want +got):\n%s", diff) + } + }) + + t.Run("primitive output", func(t *testing.T) { + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "count"}, + func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, int, error) { + return nil, 42, nil + }) + + ct, st := NewInMemoryTransports() + if _, err := server.Connect(ctx, st, nil); err != nil { + t.Fatal(err) + } + client := NewClient(testImpl, nil) + cs, err := client.Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + res, err := cs.CallTool(ctx, &CallToolParams{Name: "count"}) + if err != nil { + t.Fatal(err) + } + if res.IsError { + t.Fatalf("unexpected tool error: %v", res.Content) + } + if diff := cmp.Diff(json.Number("42"), res.StructuredContent); diff != "" { + t.Errorf("structured content mismatch (-want +got):\n%s", diff) + } + }) +} + +// TestAddToolInputSchemaComposition verifies SEP-2106 (input side): composition +// keywords such as oneOf are allowed on the input schema alongside +// type:"object". +func TestAddToolInputSchemaComposition(t *testing.T) { + server := NewServer(testImpl, nil) + tool := &Tool{ + Name: "lookup", + InputSchema: &jsonschema.Schema{ + Type: "object", + OneOf: []*jsonschema.Schema{ + {Required: []string{"id"}, Properties: map[string]*jsonschema.Schema{"id": {Type: "string"}}}, + {Required: []string{"name"}, Properties: map[string]*jsonschema.Schema{"name": {Type: "string"}}}, + }, + }, + } + // Must not panic. + server.AddTool(tool, func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{}, nil + }) +} + type schema = jsonschema.Schema func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out, wantIn, wantOut any, wantErrContaining string) { @@ -497,7 +896,7 @@ func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out th := func(context.Context, *CallToolRequest, In) (*CallToolResult, Out, error) { return nil, out, nil } - gott, goth, err := toolForErr(tool, th) + gott, goth, err := toolForErr(tool, th, nil) if err != nil { t.Fatal(err) } @@ -514,12 +913,19 @@ func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out } result, err := goth(context.Background(), ctr) if wantErrContaining != "" { - if err == nil { - t.Errorf("got nil error, want error containing %q", wantErrContaining) - } else { + // Input validation errors are returned as tool results with IsError=true, + // not as Go errors. Check both possibilities. + if err != nil { if !strings.Contains(err.Error(), wantErrContaining) { t.Errorf("got error %q, want containing %q", err, wantErrContaining) } + } else if result != nil && result.IsError { + text := result.Content[0].(*TextContent).Text + if !strings.Contains(text, wantErrContaining) { + t.Errorf("got tool error %q, want containing %q", text, wantErrContaining) + } + } else { + t.Errorf("got no error, want error containing %q", wantErrContaining) } } else if err != nil { t.Errorf("got error %v, want no error", err) @@ -535,6 +941,214 @@ func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out } } +// TestClientRootCapabilities verifies that the server correctly observes +// RootsV2 for various client capability configurations. This tests the fix +// for #607. +func TestClientRootCapabilities(t *testing.T) { + testCases := []struct { + name string + capabilities *string // JSON for the capabilities field; nil means omit the field + wantRootsV2 *RootCapabilities + }{ + { + name: "capabilities field omitted", + capabilities: nil, + wantRootsV2: nil, + }, + { + name: "empty capabilities", + capabilities: ptr(`{}`), + wantRootsV2: nil, + }, + { + name: "capabilities with no roots", + capabilities: ptr(`{"sampling": {}}`), + wantRootsV2: nil, + }, + { + name: "capabilities with empty roots", + capabilities: ptr(`{"roots": {}}`), + wantRootsV2: &RootCapabilities{ListChanged: false}, + }, + { + name: "capabilities with roots without listChanged", + capabilities: ptr(`{"roots": {"listChanged": false}}`), + wantRootsV2: &RootCapabilities{ListChanged: false}, + }, + { + name: "capabilities with roots with listChanged", + capabilities: ptr(`{"roots": {"listChanged": true}}`), + wantRootsV2: &RootCapabilities{ListChanged: true}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + + // Create a minimal server. + impl := &Implementation{Name: "testServer", Version: "v1.0.0"} + s := NewServer(impl, nil) + + // Connect the server. + cTransport, sTransport := NewInMemoryTransports() + ss, err := s.Connect(ctx, sTransport, nil) + if err != nil { + t.Fatal(err) + } + + // Connect the client JSON-RPC connection (raw, no client). + cConn, err := cTransport.Connect(ctx) + if err != nil { + t.Fatal(err) + } + + // Build initialize params, optionally including capabilities. + var initParams json.RawMessage + if tc.capabilities != nil { + initParams = json.RawMessage(`{ + "protocolVersion": "2025-11-25", + "capabilities": ` + *tc.capabilities + `, + "clientInfo": {"name": "TestClient", "version": "1.0.0"} + }`) + } else { + initParams = json.RawMessage(`{ + "protocolVersion": "2025-11-25", + "clientInfo": {"name": "TestClient", "version": "1.0.0"} + }`) + } + + initReq, err := jsonrpc2.NewCall(jsonrpc2.Int64ID(1), "initialize", initParams) + if err != nil { + t.Fatal(err) + } + + if err := cConn.Write(ctx, initReq); err != nil { + t.Fatalf("Write failed: %v", err) + } + + // Read the initialize response. + msg, err := cConn.Read(ctx) + if err != nil { + t.Fatalf("Read failed: %v", err) + } + resp, ok := msg.(*jsonrpc2.Response) + if !ok { + t.Fatalf("expected Response, got %T", msg) + } + if resp.Error != nil { + t.Fatalf("initialize failed: %v", resp.Error) + } + + // Verify that the server session has the correct RootsV2 value. + params := ss.InitializeParams() + if params == nil { + t.Fatal("InitializeParams is nil") + } + + var gotRootsV2 *RootCapabilities + if params.Capabilities != nil { + gotRootsV2 = params.Capabilities.RootsV2 + } + if diff := cmp.Diff(tc.wantRootsV2, gotRootsV2); diff != "" { + t.Errorf("RootsV2 mismatch (-want +got):\n%s", diff) + } + + // Close the client connection. + if err := cConn.Close(); err != nil { + t.Fatalf("Stream.Close failed: %v", err) + } + ss.Wait() + }) + } +} + +func TestServerRejectsDuplicateInitialize(t *testing.T) { + ctx := context.Background() + + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + cTransport, sTransport := NewInMemoryTransports() + ss, err := server.Connect(ctx, sTransport, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + cConn, err := cTransport.Connect(ctx) + if err != nil { + t.Fatal(err) + } + defer cConn.Close() + + firstParams := json.RawMessage(`{ + "protocolVersion": "2025-11-25", + "clientInfo": {"name": "first-client", "version": "1.0.0"} + }`) + firstReq, err := jsonrpc2.NewCall(jsonrpc2.Int64ID(1), methodInitialize, firstParams) + if err != nil { + t.Fatal(err) + } + if err := cConn.Write(ctx, firstReq); err != nil { + t.Fatalf("first initialize write failed: %v", err) + } + msg, err := cConn.Read(ctx) + if err != nil { + t.Fatalf("first initialize read failed: %v", err) + } + resp, ok := msg.(*jsonrpc2.Response) + if !ok { + t.Fatalf("expected Response, got %T", msg) + } + if resp.Error != nil { + t.Fatalf("first initialize failed: %v", resp.Error) + } + + initializedReq, err := jsonrpc2.NewNotification(notificationInitialized, &InitializedParams{}) + if err != nil { + t.Fatal(err) + } + if err := cConn.Write(ctx, initializedReq); err != nil { + t.Fatalf("initialized notification write failed: %v", err) + } + + secondParams := json.RawMessage(`{ + "protocolVersion": "2024-11-05", + "clientInfo": {"name": "second-client", "version": "2.0.0"} + }`) + secondReq, err := jsonrpc2.NewCall(jsonrpc2.Int64ID(2), methodInitialize, secondParams) + if err != nil { + t.Fatal(err) + } + if err := cConn.Write(ctx, secondReq); err != nil { + t.Fatalf("second initialize write failed: %v", err) + } + msg, err = cConn.Read(ctx) + if err != nil { + t.Fatalf("second initialize read failed: %v", err) + } + resp, ok = msg.(*jsonrpc2.Response) + if !ok { + t.Fatalf("expected Response, got %T", msg) + } + if resp.Error == nil { + t.Fatal("second initialize unexpectedly succeeded") + } + if !strings.Contains(resp.Error.Error(), `duplicate "initialize" received`) { + t.Fatalf("second initialize error = %v, want duplicate initialize", resp.Error) + } + + got := ss.InitializeParams() + if got == nil { + t.Fatal("InitializeParams is nil") + } + if got.ProtocolVersion != "2025-11-25" { + t.Fatalf("ProtocolVersion = %q, want first initialize value", got.ProtocolVersion) + } + if got.ClientInfo == nil || got.ClientInfo.Name != "first-client" { + t.Fatalf("ClientInfo = %#v, want first initialize value", got.ClientInfo) + } +} + // TODO: move this to tool_test.go func TestToolForSchemas(t *testing.T) { // Validate that toolForErr handles schemas properly. @@ -547,10 +1161,34 @@ func TestToolForSchemas(t *testing.T) { var ( falseSchema = &schema{Not: &schema{}} - inSchema = &schema{Type: "object", AdditionalProperties: falseSchema, Properties: map[string]*schema{"p": {Type: "integer"}}} - inSchema2 = &schema{Type: "object", AdditionalProperties: falseSchema, Properties: map[string]*schema{"p": {Type: "string"}}} - outSchema = &schema{Type: "object", AdditionalProperties: falseSchema, Properties: map[string]*schema{"b": {Type: "boolean"}}} - outSchema2 = &schema{Type: "object", AdditionalProperties: falseSchema, Properties: map[string]*schema{"b": {Type: "integer"}}} + inSchema = &schema{ + Type: "object", + AdditionalProperties: falseSchema, + Properties: map[string]*schema{"p": {Type: "integer"}}, + PropertyOrder: []string{"p"}, + } + inSchema2 = &schema{ + Type: "object", + AdditionalProperties: falseSchema, + Properties: map[string]*schema{"p": {Type: "string"}}, + } + inSchema3 = &schema{ + Type: "object", + AdditionalProperties: falseSchema, + Properties: map[string]*schema{}, // empty map is preserved + } + outSchema = &schema{ + Type: "object", + AdditionalProperties: falseSchema, + Properties: map[string]*schema{"b": {Type: "boolean"}}, + PropertyOrder: []string{"b"}, + } + outSchema2 = &schema{ + Type: "object", + AdditionalProperties: falseSchema, + Properties: map[string]*schema{"b": {Type: "integer"}}, + PropertyOrder: []string{"b"}, + } ) // Infer both schemas. @@ -564,6 +1202,8 @@ func TestToolForSchemas(t *testing.T) { testToolForSchema[in, any](t, &Tool{}, `{"p":"x"}`, 0, inSchema, nil, `want "integer"`) // Tool sets input schema: that is what's used. testToolForSchema[in, any](t, &Tool{InputSchema: inSchema2}, `{"p":3}`, 0, inSchema2, nil, `want "string"`) + // Tool sets input schema, empty properties map. + testToolForSchema[in, any](t, &Tool{InputSchema: inSchema3}, `{}`, 0, inSchema3, nil, "") // Tool sets output schema: that is what's used, and validation happens. testToolForSchema[in, any](t, &Tool{OutputSchema: outSchema2}, `{"p":3}`, out{true}, inSchema, outSchema2, `want "integer"`) @@ -585,6 +1225,1189 @@ func TestToolForSchemas(t *testing.T) { "AsOf": {Type: "string"}, "Source": {Type: "string"}, }, + PropertyOrder: []string{"Summary", "AsOf", "Source"}, }, "") } + +// TestServerCapabilitiesOverWire verifies that server capabilities are +// correctly sent over the wire during initialization. +func TestServerCapabilitiesOverWire(t *testing.T) { + tool := &Tool{Name: "test-tool", InputSchema: &jsonschema.Schema{Type: "object"}} + + testCases := []struct { + name string + serverOpts *ServerOptions + configureServer func(s *Server) + wantCapabilities *ServerCapabilities + }{ + { + name: "Default capabilities", + serverOpts: nil, + configureServer: func(s *Server) {}, + wantCapabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + }, + }, + { + name: "Custom Capabilities with tools", + serverOpts: &ServerOptions{ + Capabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: false}, + }, + }, + configureServer: func(s *Server) {}, + wantCapabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: false}, + }, + }, + { + name: "Dynamic tool capability", + serverOpts: &ServerOptions{ + Capabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + }, + }, + configureServer: func(s *Server) { + s.AddTool(tool, nil) + }, + wantCapabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Tools: &ToolCapabilities{ListChanged: true}, + }, + }, + { + name: "Extensions over wire", + serverOpts: func() *ServerOptions { + caps := &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + } + caps.AddExtension("io.example/ext", map[string]any{"key": "value"}) + return &ServerOptions{Capabilities: caps} + }(), + configureServer: func(s *Server) {}, + wantCapabilities: &ServerCapabilities{ + Extensions: map[string]any{ + "io.example/ext": map[string]any{"key": "value"}, + }, + Logging: &LoggingCapabilities{}, + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + + // Create server. + impl := &Implementation{Name: "testServer", Version: "v1.0.0"} + server := NewServer(impl, tc.serverOpts) + tc.configureServer(server) + + // Connect client and server. + cTransport, sTransport := NewInMemoryTransports() + ss, err := server.Connect(ctx, sTransport, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, nil) + cs, err := client.Connect(ctx, cTransport, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Check that the client received the expected capabilities. + initResult := cs.InitializeResult() + if initResult == nil { + t.Fatal("InitializeResult is nil") + } + + if diff := cmp.Diff(tc.wantCapabilities, initResult.Capabilities); diff != "" { + t.Errorf("Capabilities mismatch (-want +got):\n%s", diff) + } + }) + } +} + +// SEP-2575 removes the initialization handshake. An `initialize` request +// that opts into the new protocol via `_meta.protocolVersion` must be +// rejected with `Method not found` (-32601). +func TestServerSessionHandle_RejectsInitializeOnNewProtocol(t *testing.T) { + + tests := []struct { + name string + params any + wantReject bool + }{ + { + name: "initialize with new-protocol _meta is rejected", + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "c", "version": "1"}, + MetaKeyClientCapabilities: map[string]any{}, + }, + "protocolVersion": protocolVersion20260728, + }, + wantReject: true, + }, + { + name: "initialize without _meta is allowed (old protocol)", + params: map[string]any{ + "protocolVersion": protocolVersion20251125, + }, + wantReject: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ss := &ServerSession{server: NewServer(testImpl, nil)} + id, err := jsonrpc.MakeID("test") + if err != nil { + t.Fatal(err) + } + req := &jsonrpc.Request{ + ID: id, + Method: methodInitialize, + Params: mustMarshal(tc.params), + } + _, err = ss.handle(context.Background(), req) + if tc.wantReject { + if err == nil { + t.Fatal("expected error rejecting initialize, got nil") + } + var jerr *jsonrpc.Error + if !errors.As(err, &jerr) { + t.Fatalf("error type = %T, want *jsonrpc.Error so the wire returns the right code", err) + } + if jerr.Code != jsonrpc.CodeMethodNotFound { + t.Errorf("error code = %d, want %d (CodeMethodNotFound = -32601)", jerr.Code, jsonrpc.CodeMethodNotFound) + } + if !strings.Contains(jerr.Message, "initialize") { + t.Errorf("error message %q does not mention %q", jerr.Message, "initialize") + } + } else { + // Old-protocol initialize should be dispatched normally; any + // CodeMethodNotFound here would mean the rejection branch + // fired incorrectly. + var jerr *jsonrpc.Error + if errors.As(err, &jerr) && jerr.Code == jsonrpc.CodeMethodNotFound { + t.Errorf("old-protocol initialize was incorrectly rejected: %v", err) + } + } + }) + } + + t.Run("rejection error encodes to wire as code -32601", func(t *testing.T) { + // Belt-and-braces check that the error type produced by handle() + // actually serializes to JSON-RPC code -32601, not a bare 0. + ss := &ServerSession{server: NewServer(testImpl, nil)} + id, err := jsonrpc.MakeID("test") + if err != nil { + t.Fatal(err) + } + req := &jsonrpc.Request{ + ID: id, + Method: methodInitialize, + Params: mustMarshal(map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "c", "version": "1"}, + MetaKeyClientCapabilities: map[string]any{}, + }, + "protocolVersion": protocolVersion20260728, + }), + } + _, handleErr := ss.handle(context.Background(), req) + if handleErr == nil { + t.Fatal("expected rejection error, got nil") + } + data, encErr := jsonrpc.EncodeMessage(&jsonrpc.Response{ID: id, Error: handleErr.(*jsonrpc.Error)}) + if encErr != nil { + t.Fatal(encErr) + } + var wire struct { + Error struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(data, &wire); err != nil { + t.Fatal(err) + } + if wire.Error.Code != jsonrpc.CodeMethodNotFound { + t.Errorf("wire error code = %d, want %d; full response = %s", wire.Error.Code, jsonrpc.CodeMethodNotFound, data) + } + }) +} + +// newProtocolParams merges fields into the _meta that opts a request into the +// sessionless protocol (SEP-2575). +func newProtocolParams(fields map[string]any) map[string]any { + params := map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "c", "version": "1"}, + MetaKeyClientCapabilities: map[string]any{}, + }, + } + for k, v := range fields { + params[k] = v + } + return params +} + +func TestServerSessionHandle_SetsResultTypeOnNewProtocol(t *testing.T) { + server := NewServer(testImpl, &ServerOptions{ + CompletionHandler: func(context.Context, *CompleteRequest) (*CompleteResult, error) { + return &CompleteResult{ + Completion: CompletionResultDetails{ + Values: []string{"go"}, + }, + }, nil + }, + }) + AddTool(server, &Tool{Name: "tool"}, func(context.Context, *CallToolRequest, struct{}) (*CallToolResult, any, error) { + return &CallToolResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Continue?"}}, + }, nil, nil + }) + server.AddPrompt(&Prompt{Name: "prompt"}, func(context.Context, *GetPromptRequest) (*GetPromptResult, error) { + return &GetPromptResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Continue?"}}, + }, nil + }) + server.AddResource(&Resource{URI: "test://resource", Name: "resource"}, func(context.Context, *ReadResourceRequest) (*ReadResourceResult, error) { + return &ReadResourceResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Continue?"}}, + }, nil + }) + AddTool(server, &Tool{Name: "toolComplete"}, func(context.Context, *CallToolRequest, struct{}) (*CallToolResult, any, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + server.AddPrompt(&Prompt{Name: "promptComplete"}, func(context.Context, *GetPromptRequest) (*GetPromptResult, error) { + return &GetPromptResult{ + Messages: []*PromptMessage{{Role: "assistant", Content: &TextContent{Text: "ok"}}}, + }, nil + }) + server.AddResource(&Resource{URI: "test://resource-complete", Name: "resourceComplete"}, func(context.Context, *ReadResourceRequest) (*ReadResourceResult, error) { + return &ReadResourceResult{ + Contents: []*ResourceContents{{URI: "test://resource-complete", Text: "ok"}}, + }, nil + }) + + tests := []struct { + name string + method string + params map[string]any + want resultType + }{ + { + name: "discover", + method: methodDiscover, + params: newProtocolParams(nil), + want: resultTypeComplete, + }, + { + name: "tools list", + method: methodListTools, + params: newProtocolParams(nil), + want: resultTypeComplete, + }, + { + name: "prompts list", + method: methodListPrompts, + params: newProtocolParams(nil), + want: resultTypeComplete, + }, + { + name: "resources list", + method: methodListResources, + params: newProtocolParams(nil), + want: resultTypeComplete, + }, + { + name: "resource templates list", + method: methodListResourceTemplates, + params: newProtocolParams(nil), + want: resultTypeComplete, + }, + { + name: "complete", + method: methodComplete, + params: newProtocolParams(map[string]any{ + "argument": map[string]any{ + "name": "language", + "value": "g", + }, + "ref": map[string]any{ + "type": "ref/prompt", + "name": "code_review", + }, + }), + want: resultTypeComplete, + }, + { + name: "tool input required", + method: methodCallTool, + params: newProtocolParams(map[string]any{"name": "tool", "arguments": map[string]any{}}), + want: resultTypeInputRequired, + }, + { + name: "prompt input required", + method: methodGetPrompt, + params: newProtocolParams(map[string]any{"name": "prompt"}), + want: resultTypeInputRequired, + }, + { + name: "resource input required", + method: methodReadResource, + params: newProtocolParams(map[string]any{"uri": "test://resource"}), + want: resultTypeInputRequired, + }, + { + name: "tool complete", + method: methodCallTool, + params: newProtocolParams(map[string]any{"name": "toolComplete", "arguments": map[string]any{}}), + want: resultTypeComplete, + }, + { + name: "prompt complete", + method: methodGetPrompt, + params: newProtocolParams(map[string]any{"name": "promptComplete"}), + want: resultTypeComplete, + }, + { + name: "resource complete", + method: methodReadResource, + params: newProtocolParams(map[string]any{"uri": "test://resource-complete"}), + want: resultTypeComplete, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := resultTypeOf(t, server, nil, tc.method, tc.params); got != string(tc.want) { + t.Errorf("resultType = %q, want %q", got, tc.want) + } + }) + } +} + +// TestServerSessionHandle_SetsResultTypeWhenMiddlewareShortCircuits is a +// regression test for #1225. User middleware wraps outside +// serverMultiRoundTripMiddleware, so a middleware that returns a result +// without calling next bypasses both that shim and the dispatcher. +func TestServerSessionHandle_SetsResultTypeWhenMiddlewareShortCircuits(t *testing.T) { + tests := []struct { + name string + method string + params map[string]any + result Result + want resultType + }{ + { + name: "tools/call", + method: methodCallTool, + params: newProtocolParams(map[string]any{"name": "tool", "arguments": map[string]any{}}), + result: &CallToolResult{Content: []Content{&TextContent{Text: "denied"}}}, + want: resultTypeComplete, + }, + { + name: "prompts/get", + method: methodGetPrompt, + params: newProtocolParams(map[string]any{"name": "prompt"}), + result: &GetPromptResult{ + Messages: []*PromptMessage{{Role: "assistant", Content: &TextContent{Text: "denied"}}}, + }, + want: resultTypeComplete, + }, + { + name: "resources/read", + method: methodReadResource, + params: newProtocolParams(map[string]any{"uri": "test://resource"}), + result: &ReadResourceResult{ + Contents: []*ResourceContents{{URI: "test://resource", Text: "denied"}}, + }, + want: resultTypeComplete, + }, + { + // setResultType is unexported, so populating InputRequests is + // the only way a middleware can ask for input. + name: "tools/call input required", + method: methodCallTool, + params: newProtocolParams(map[string]any{"name": "tool", "arguments": map[string]any{}}), + result: &CallToolResult{ + InputRequests: InputRequestMap{"confirm": &ElicitParams{Message: "Continue?"}}, + }, + want: resultTypeInputRequired, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + // Nothing is registered: if the middleware fell through to next, + // dispatch would fail rather than return a result. + server := NewServer(testImpl, nil) + server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == tc.method { + return tc.result, nil + } + return next(ctx, method, req) + } + }) + if got := resultTypeOf(t, server, nil, tc.method, tc.params); got != string(tc.want) { + t.Errorf("resultType = %q, want %q", got, tc.want) + } + }) + } +} + +// TestServerSessionHandle_ResultTypeGate pins when a result carries resultType. +// annotateResultType is the only setter, and handle gates it on the request's +// _meta rather than on the version the session negotiated. So the last case +// below gets no resultType even though the session speaks 2026-07-28. Only a +// client that mixes protocol versions reaches that, so it is recorded rather +// than fixed. +func TestServerSessionHandle_ResultTypeGate(t *testing.T) { + version := func(v string) func(*ServerSessionState) { + return func(s *ServerSessionState) { + s.InitializeParams = &InitializeParams{ProtocolVersion: v} + } + } + callParams := map[string]any{"name": "tool", "arguments": map[string]any{}} + + tests := []struct { + name string + state func(*ServerSessionState) + params map[string]any + want resultType // "" means the field must be absent + }{ + { + name: "legacy session, no _meta", + state: version(protocolVersion20251125), + params: callParams, + }, + { + name: "legacy session, new-protocol _meta", + state: version(protocolVersion20251125), + params: newProtocolParams(callParams), + want: resultTypeComplete, + }, + { + name: "new-protocol session, no _meta", + state: version(protocolVersion20260728), + params: callParams, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + server := NewServer(testImpl, nil) + // Short-circuit so this measures handle's gate, not the + // dispatcher's. + server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(context.Context, string, Request) (Result, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil + } + }) + if got := resultTypeOf(t, server, tc.state, methodCallTool, tc.params); got != string(tc.want) { + t.Errorf("resultType = %q, want %q", got, tc.want) + } + }) + } +} + +// TestServerSessionHandle_NilResultFromMiddleware checks that a middleware +// returning no result does not panic. A typed nil satisfies Result, so the +// annotation helpers would otherwise dereference it. +func TestServerSessionHandle_NilResultFromMiddleware(t *testing.T) { + tests := []struct { + name string + result Result + }{ + {"nil interface", nil}, + {"typed nil", (*CallToolResult)(nil)}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + server := NewServer(testImpl, nil) + server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(context.Context, string, Request) (Result, error) { + return tc.result, nil + } + }) + ss := &ServerSession{server: server} + id, err := jsonrpc.MakeID("test") + if err != nil { + t.Fatal(err) + } + res, err := ss.handle(context.Background(), &jsonrpc.Request{ + ID: id, + Method: methodCallTool, + Params: mustMarshal(newProtocolParams(map[string]any{"name": "tool", "arguments": map[string]any{}})), + }) + if err != nil { + t.Fatalf("handle: %v", err) + } + data, err := json.Marshal(res) + if err != nil { + t.Fatal(err) + } + if got := string(data); got != "null" { + t.Errorf("result = %s, want null", got) + } + }) + } +} + +// resultTypeOf handles one request on a transportless session and reports the +// resultType of the marshaled result, or "" if the field is absent. +func resultTypeOf(t *testing.T, server *Server, state func(*ServerSessionState), method string, params map[string]any) string { + t.Helper() + ss := &ServerSession{server: server} + if state != nil { + ss.updateState(state) + } + id, err := jsonrpc.MakeID("test") + if err != nil { + t.Fatal(err) + } + result, err := ss.handle(context.Background(), &jsonrpc.Request{ + ID: id, + Method: method, + Params: mustMarshal(params), + }) + if err != nil { + t.Fatal(err) + } + data, err := json.Marshal(result) + if err != nil { + t.Fatal(err) + } + var fields map[string]json.RawMessage + if err := json.Unmarshal(data, &fields); err != nil { + t.Fatal(err) + } + raw, ok := fields["resultType"] + if !ok { + return "" + } + var got string + if err := json.Unmarshal(raw, &got); err != nil { + t.Fatalf("unmarshaling resultType from %s: %v", data, err) + } + return got +} + +// TestServerSessionHandle_RejectsRemovedMethodsOnNewProtocol verifies that +// the methods removed by SEP-2575 (`initialize`, `notifications/initialized`, +// `ping`) all return Method not found when the request opts into the new +// protocol via `_meta.protocolVersion`. +func TestServerSessionHandle_RejectsRemovedMethodsOnNewProtocol(t *testing.T) { + + newProtoMeta := map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "c", "version": "1"}, + MetaKeyClientCapabilities: map[string]any{}, + }, + } + + tests := []struct { + name string + method string + }{ + {"initialize", methodInitialize}, + {"ping", methodPing}, + {"notifications/initialized", notificationInitialized}, + } + + for _, tc := range tests { + t.Run(tc.name+" rejected on new protocol", func(t *testing.T) { + ss := &ServerSession{server: NewServer(testImpl, nil)} + id, err := jsonrpc.MakeID("test") + if err != nil { + t.Fatal(err) + } + req := &jsonrpc.Request{ + ID: id, + Method: tc.method, + Params: mustMarshal(newProtoMeta), + } + _, err = ss.handle(context.Background(), req) + if err == nil { + t.Fatalf("method %q on new protocol: got nil error, want CodeMethodNotFound", tc.method) + } + var jerr *jsonrpc.Error + if !errors.As(err, &jerr) { + t.Fatalf("error type = %T, want *jsonrpc.Error", err) + } + if jerr.Code != jsonrpc.CodeMethodNotFound { + t.Errorf("method %q: code = %d, want %d", tc.method, jerr.Code, jsonrpc.CodeMethodNotFound) + } + if !strings.Contains(jerr.Message, tc.method) { + t.Errorf("method %q: message %q does not mention method name", tc.method, jerr.Message) + } + }) + } +} + +// TestServerSession_RejectsServerInitiatedRequests verifies that +// SEP-2322 / SEP-2575 is enforced at the API surface: [ServerSession.Elicit], +// [ServerSession.CreateMessage], [ServerSession.CreateMessageWithTools], and +// [ServerSession.ListRoots] must refuse to send a server-to-client request +// when the session is negotiated at protocol version >= 2026-07-28, and must +// remain functional on pre-2026-07-28 sessions. +func TestServerSession_RejectsServerInitiated(t *testing.T) { + ctx := context.Background() + + tests := []struct { + name string + call func(*ServerSession) error + }{ + { + name: "Elicit", + call: func(ss *ServerSession) error { + _, err := ss.Elicit(ctx, &ElicitParams{Message: "hi"}) + return err + }, + }, + { + name: "CreateMessage", + call: func(ss *ServerSession) error { + _, err := ss.CreateMessage(ctx, &CreateMessageParams{}) + return err + }, + }, + { + name: "CreateMessageWithTools", + call: func(ss *ServerSession) error { + _, err := ss.CreateMessageWithTools(ctx, &CreateMessageWithToolsParams{}) + return err + }, + }, + { + name: "ListRoots", + call: func(ss *ServerSession) error { + _, err := ss.ListRoots(ctx, nil) + return err + }, + }, + } + + // Cover both branches of the era gate: modern must reject, legacy must let + // the call proceed to the wire (where it either succeeds or fails on + // something unrelated to the guard, such as missing client capabilities). + for _, protoVer := range []string{protocolVersion20260728, protocolVersion20251125} { + for _, tc := range tests { + t.Run(fmt.Sprintf("%s/%s", protoVer, tc.name), func(t *testing.T) { + ct, st := NewInMemoryTransports() + s := NewServer(testImpl, nil) + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ss.Close() }) + + // Client advertises every server-to-client capability so a + // legacy-era call is not rejected for capability reasons + // instead of era reasons. + c := NewClient(testImpl, &ClientOptions{ + ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) { + return &ElicitResult{Action: "cancel"}, nil + }, + CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) { + return &CreateMessageResult{Model: "m", Role: "assistant", Content: &TextContent{Text: "ok"}}, nil + }, + }) + c.AddRoots(&Root{URI: "file:///workspace"}) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protoVer}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = cs.Close() }) + + gotErr := tc.call(ss) + modern := protoVer >= protocolVersion20260728 + if modern { + if gotErr == nil { + t.Fatalf("%s on %s: got nil error, want era-gate rejection", tc.name, protoVer) + } + wantSubstr := "cannot be sent while serving a request on protocol version " + protoVer + if !strings.Contains(gotErr.Error(), wantSubstr) { + t.Errorf("%s on %s: error %q does not contain %q", tc.name, protoVer, gotErr.Error(), wantSubstr) + } + if !strings.Contains(gotErr.Error(), "multi round-trip requests") { + t.Errorf("%s on %s: error %q does not steer to MRTR", tc.name, protoVer, gotErr.Error()) + } + } else { + if gotErr != nil { + t.Errorf("%s on %s: got error %v, want nil (era gate must not fire on legacy)", tc.name, protoVer, gotErr) + } + } + }) + } + } +} + +func TestServerSetCacheable(t *testing.T) { + ctx := context.Background() + // The values an opinionated resource handler sets on its own result. + handlerValue := Cacheable{TTLMs: 9_000, CacheScope: "private"} + + for _, tc := range []struct { + name string + callback func(context.Context, Request, *Cacheable) + wantSDK Cacheable // for a result whose fields nobody set + wantHandler Cacheable // for a result the handler set fields on + }{ + { + name: "no callback", + wantSDK: Cacheable{TTLMs: 0, CacheScope: "public"}, + wantHandler: handlerValue, + }, + { + name: "callback overrules the handler", + callback: func(_ context.Context, _ Request, c *Cacheable) { + c.TTLMs, c.CacheScope = 60_000, "private" + }, + wantSDK: Cacheable{TTLMs: 60_000, CacheScope: "private"}, + wantHandler: Cacheable{TTLMs: 60_000, CacheScope: "private"}, + }, + { + // The callback is handed what the handler produced, so it can + // single out the results that expressed no opinion. The scope it + // leaves empty is still filled with the protocol default. + name: "callback defers to the handler", + callback: func(_ context.Context, _ Request, c *Cacheable) { + if c.CacheScope == "" { + c.TTLMs = 1_000 + } + }, + wantSDK: Cacheable{TTLMs: 1_000, CacheScope: "public"}, + wantHandler: handlerValue, + }, + } { + t.Run(tc.name, func(t *testing.T) { + s := NewServer(testImpl, &ServerOptions{SetCacheable: tc.callback}) + s.AddResource(&Resource{URI: "test://plain", Name: "plain"}, + func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + return &ReadResourceResult{Contents: []*ResourceContents{{URI: req.Params.URI, Text: "x"}}}, nil + }) + s.AddResource(&Resource{URI: "test://opinionated", Name: "opinionated"}, + func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + return &ReadResourceResult{ + Cacheable: handlerValue, + Contents: []*ResourceContents{{URI: req.Params.URI, Text: "x"}}, + }, nil + }) + cs := mustConnect(t, s, nil) + + check := func(label string, got, want Cacheable) { + t.Helper() + if got != want { + t.Errorf("%s Cacheable = %+v, want %+v", label, got, want) + } + } + + tools, err := cs.ListTools(ctx, nil) // built by the SDK, never by a handler + if err != nil { + t.Fatal(err) + } + check("tools/list", tools.Cacheable, tc.wantSDK) + + for uri, want := range map[string]Cacheable{ + "test://plain": tc.wantSDK, + "test://opinionated": tc.wantHandler, + } { + res, err := cs.ReadResource(ctx, &ReadResourceParams{URI: uri}) + if err != nil { + t.Fatal(err) + } + check(uri, res.Cacheable, want) + } + }) + } +} + +// TestServerSupportedProtocolVersions verifies that +// [ServerOptions.SupportedProtocolVersions] narrows both the versions the +// server advertises in server/discover and the versions it negotiates, +// and that it cannot widen SDK support. +func TestServerSupportedProtocolVersions(t *testing.T) { + ctx := context.Background() + + t.Run("discover advertises the configured subset", func(t *testing.T) { + server := NewServer(testImpl, &ServerOptions{ + // Deliberately unordered: the server must advertise newest first. + SupportedProtocolVersions: []string{protocolVersion20251125, protocolVersion20260728}, + }) + var advertised []string + server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + res, err := next(ctx, method, req) + if method == methodDiscover && err == nil { + advertised = res.(*DiscoverResult).SupportedVersions + } + return res, err + } + }) + + ct, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server.Connect: %v", err) + } + defer ss.Close() + + cs, err := NewClient(testImpl, nil).Connect(ctx, ct, &ClientSessionOptions{ + ProtocolVersion: protocolVersion20260728, + }) + if err != nil { + t.Fatalf("client.Connect: %v", err) + } + defer cs.Close() + + want := []string{protocolVersion20260728, protocolVersion20251125} + if diff := cmp.Diff(want, advertised); diff != "" { + t.Errorf("DiscoverResult.SupportedVersions mismatch (-want +got):\n%s", diff) + } + }) + + t.Run("initialize negotiates within the configured subset", func(t *testing.T) { + tests := []struct { + name string + supported []string + client string + want string + }{ + { + name: "requested version supported", + supported: []string{protocolVersion20251125, protocolVersion20250618}, + client: protocolVersion20250618, + want: protocolVersion20250618, + }, + { + name: "requested version excluded", + supported: []string{protocolVersion20251125, protocolVersion20250618}, + client: protocolVersion20241105, + want: protocolVersion20251125, + }, + { + name: "fallback is the newest configured version", + supported: []string{protocolVersion20250326}, + client: protocolVersion20241105, + want: protocolVersion20250326, + }, + { + // Nothing the server supports can carry an initialize + // handshake, so the answer names a handshake-era version the + // client can act on rather than one it could mistake for the + // new protocol. + name: "server with no handshake version", + supported: []string{protocolVersion20260728}, + client: protocolVersion20250618, + want: protocolVersion20251125, + }, + { + name: "unrestricted server keeps the default behavior", + supported: nil, + client: protocolVersion20241105, + want: protocolVersion20241105, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := NewServer(testImpl, &ServerOptions{SupportedProtocolVersions: test.supported}) + ct, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server.Connect: %v", err) + } + defer ss.Close() + + cs, err := NewClient(testImpl, nil).Connect(ctx, ct, &ClientSessionOptions{ + ProtocolVersion: test.client, + }) + if err != nil { + t.Fatalf("client.Connect: %v", err) + } + defer cs.Close() + + if got := cs.InitializeResult().ProtocolVersion; got != test.want { + t.Errorf("negotiated protocol version = %q, want %q", got, test.want) + } + }) + } + }) + + t.Run("unknown version panics", func(t *testing.T) { + defer func() { + if recover() == nil { + t.Error("NewServer did not panic on an unsupported protocol version") + } + }() + NewServer(testImpl, &ServerOptions{SupportedProtocolVersions: []string{"1999-01-01"}}) + }) +} + +// TestServerSupportedProtocolVersions_NewProtocol verifies that a request +// using the SEP-2575 per-request protocol is rejected with +// [CodeUnsupportedProtocolVersion] when its version is excluded by +// [ServerOptions.SupportedProtocolVersions]. +func TestServerSupportedProtocolVersions_NewProtocol(t *testing.T) { + ctx := context.Background() + server := NewServer(testImpl, &ServerOptions{ + SupportedProtocolVersions: []string{protocolVersion20251125}, + }) + ct, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server.Connect: %v", err) + } + defer ss.Close() + _ = ct + + params := fmt.Sprintf(`{"_meta":{%q:%q,%q:{}}}`, + MetaKeyProtocolVersion, protocolVersion20260728, MetaKeyClientCapabilities) + _, err = ss.handle(ctx, &jsonrpc.Request{ + ID: jsonrpc2.Int64ID(1), + Method: methodListTools, + Params: json.RawMessage(params), + }) + var jerr *jsonrpc.Error + if !errors.As(err, &jerr) { + t.Fatalf("handle returned %v, want a *jsonrpc.Error", err) + } + if jerr.Code != CodeUnsupportedProtocolVersion { + t.Fatalf("error code = %d, want %d", jerr.Code, CodeUnsupportedProtocolVersion) + } + var data UnsupportedProtocolVersionData + if err := json.Unmarshal(jerr.Data, &data); err != nil { + t.Fatalf("unmarshal error data: %v", err) + } + if diff := cmp.Diff([]string{protocolVersion20251125}, data.Supported); diff != "" { + t.Errorf("UnsupportedProtocolVersionData.Supported mismatch (-want +got):\n%s", diff) + } +} + +// TestServerHandle_LegacyCallBeforeInitialize asserts that a call carrying no +// new-protocol `_meta` is refused on a session that has not completed +// 'initialize', whichever branch of the method switch it lands in. The +// handshake itself is what ends that state, and the lifecycle spec exempts +// pings, so those two are served. +// +// The gate used to sit in the default branch, so 'logging/setLevel', +// 'resources/subscribe' and 'resources/unsubscribe', which share a case that +// only refuses new-protocol requests, were served before the handshake. For +// subscribe that registered a session that never handshook as a subscriber, +// so the subtest also asserts that the refusal leaves no subscriber and runs +// no handler. +func TestServerHandle_LegacyCallBeforeInitialize(t *testing.T) { + ctx := context.Background() + const uri = "test://resource" + + tests := []struct { + name string + method string + params string + served bool + }{ + { + name: "initialize", + method: methodInitialize, + params: fmt.Sprintf(`{"protocolVersion":%q,"capabilities":{},"clientInfo":{"name":"c","version":"1"}}`, protocolVersion20251125), + served: true, + }, + {name: "ping", method: methodPing, params: `{}`, served: true}, + {name: "tools/list", method: methodListTools, params: `{}`}, + {name: "logging/setLevel", method: methodSetLevel, params: `{"level":"info"}`}, + {name: "resources/subscribe", method: methodSubscribe, params: fmt.Sprintf(`{"uri":%q}`, uri)}, + {name: "resources/unsubscribe", method: methodUnsubscribe, params: fmt.Sprintf(`{"uri":%q}`, uri)}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var handlerRan bool + server := NewServer(testImpl, &ServerOptions{ + SubscribeHandler: func(context.Context, *SubscribeRequest) error { + handlerRan = true + return nil + }, + UnsubscribeHandler: func(context.Context, *UnsubscribeRequest) error { + handlerRan = true + return nil + }, + }) + _, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server.Connect: %v", err) + } + defer ss.Close() + + res, err := ss.handle(ctx, &jsonrpc.Request{ + ID: jsonrpc2.Int64ID(1), + Method: tc.method, + Params: json.RawMessage(tc.params), + }) + if tc.served { + if err != nil { + t.Fatalf("handle(%q) before initialize returned %v, want it served", tc.method, err) + } + if res == nil { + t.Fatalf("handle(%q) before initialize returned no result", tc.method) + } + if tc.method == methodInitialize && ss.InitializeParams() == nil { + t.Errorf("initialize was answered without recording InitializeParams") + } + return + } + if err == nil { + t.Fatalf("handle(%q) served the call before initialize, want a refusal", tc.method) + } + if want := fmt.Sprintf("method %q is invalid during session initialization", tc.method); err.Error() != want { + t.Errorf("handle(%q) error = %q, want %q", tc.method, err.Error(), want) + } + if handlerRan { + t.Errorf("handle(%q) ran the server's handler for a call it refused", tc.method) + } + server.mu.Lock() + _, subscribed := server.resourceSubscriptions[uri][ss] + server.mu.Unlock() + if subscribed { + t.Errorf("handle(%q) left the session subscribed to %q", tc.method, uri) + } + }) + } +} + +// TestServerHandle_NewProtocolCallWithoutInitialize asserts that the gate +// stays off the SEP-2575 path. A fresh session that never sends 'initialize' +// is served a new-protocol 'tools/list', and the methods the gate now covers +// are answered, when a new-protocol request names them, with the error the +// new protocol already gave them (the method does not exist there), and not +// with the initialization refusal. +func TestServerHandle_NewProtocolCallWithoutInitialize(t *testing.T) { + ctx := context.Background() + const uri = "test://resource" + + tests := []struct { + name string + method string + params map[string]any + served bool + }{ + {name: "tools/list", method: methodListTools, params: newProtocolParams(nil), served: true}, + {name: "logging/setLevel", method: methodSetLevel, params: newProtocolParams(map[string]any{"level": "info"})}, + {name: "resources/subscribe", method: methodSubscribe, params: newProtocolParams(map[string]any{"uri": uri})}, + {name: "resources/unsubscribe", method: methodUnsubscribe, params: newProtocolParams(map[string]any{"uri": uri})}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + server := NewServer(testImpl, &ServerOptions{ + SubscribeHandler: func(context.Context, *SubscribeRequest) error { return nil }, + UnsubscribeHandler: func(context.Context, *UnsubscribeRequest) error { return nil }, + }) + _, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server.Connect: %v", err) + } + defer ss.Close() + + params, err := json.Marshal(tc.params) + if err != nil { + t.Fatalf("marshalling params: %v", err) + } + res, err := ss.handle(ctx, &jsonrpc.Request{ + ID: jsonrpc2.Int64ID(1), + Method: tc.method, + Params: params, + }) + if tc.served { + if err != nil { + t.Fatalf("handle(%q) on a new-protocol session returned %v, want it served", tc.method, err) + } + if res == nil { + t.Fatalf("handle(%q) on a new-protocol session returned no result", tc.method) + } + return + } + var jerr *jsonrpc.Error + if !errors.As(err, &jerr) { + t.Fatalf("handle(%q) returned %v, want a *jsonrpc.Error", tc.method, err) + } + if jerr.Code != jsonrpc.CodeMethodNotFound { + t.Errorf("handle(%q) error code = %d, want %d (removed in the new protocol)", tc.method, jerr.Code, jsonrpc.CodeMethodNotFound) + } + if strings.Contains(err.Error(), "invalid during session initialization") { + t.Errorf("handle(%q) applied the initialization gate to a new-protocol request: %v", tc.method, err) + } + }) + } +} + +// TestServerUnknownProtocolVersion_NewProtocol verifies that a request whose +// `_meta.protocolVersion` names a version the SDK does not know is rejected +// with [CodeUnsupportedProtocolVersion], and not served as a legacy handshake +// because its version string happens to sort below 2026-07-28. +func TestServerUnknownProtocolVersion_NewProtocol(t *testing.T) { + ctx := context.Background() + // Each case names a version that is not in [supportedProtocolVersions] and + // sorts below 2026-07-28. + tests := []struct { + name string + version string + extraMeta string // further `_meta` entries, each with a leading comma + }{ + // The versioning spec uses 1900-01-01 to illustrate this error. + {"documented example", "1900-01-01", fmt.Sprintf(`,%q:{}`, MetaKeyClientCapabilities)}, + {"unreleased revision", "2025-11-24", fmt.Sprintf(`,%q:{}`, MetaKeyClientCapabilities)}, + {"not a date", "1.0", fmt.Sprintf(`,%q:{}`, MetaKeyClientCapabilities)}, + // The rest of the `_meta` triple is defined by the revision the request + // names, so its absence must not mask the version itself. + {"no clientCapabilities", "1900-01-01", ""}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + server := NewServer(testImpl, nil) + _, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server.Connect: %v", err) + } + defer ss.Close() + + params := fmt.Sprintf(`{"_meta":{%q:%q%s}}`, MetaKeyProtocolVersion, tc.version, tc.extraMeta) + _, err = ss.handle(ctx, &jsonrpc.Request{ + ID: jsonrpc2.Int64ID(1), + Method: methodListTools, + Params: json.RawMessage(params), + }) + var jerr *jsonrpc.Error + if !errors.As(err, &jerr) { + t.Fatalf("handle returned %v, want a *jsonrpc.Error", err) + } + if jerr.Code != CodeUnsupportedProtocolVersion { + t.Fatalf("error code = %d, want %d", jerr.Code, CodeUnsupportedProtocolVersion) + } + var data UnsupportedProtocolVersionData + if err := json.Unmarshal(jerr.Data, &data); err != nil { + t.Fatalf("unmarshal error data: %v", err) + } + if data.Requested != tc.version { + t.Errorf("UnsupportedProtocolVersionData.Requested = %q, want %q", data.Requested, tc.version) + } + if diff := cmp.Diff(SupportedProtocolVersions(), data.Supported); diff != "" { + t.Errorf("UnsupportedProtocolVersionData.Supported mismatch (-want +got):\n%s", diff) + } + }) + } +} diff --git a/mcp/session.go b/mcp/session.go index dcf9888..08eccaf 100644 --- a/mcp/session.go +++ b/mcp/session.go @@ -22,6 +22,13 @@ type ServerSessionState struct { // InitializedParams are the parameters from 'notifications/initialized'. InitializedParams *InitializedParams `json:"initializedParams"` + // NegotiatedProtocolVersion is the protocol version agreed during + // 'initialize', which may differ from the version the client requested if + // the server does not support it. + // + // It is empty for sessions that never ran the initialize handshake. + NegotiatedProtocolVersion string `json:"negotiatedProtocolVersion,omitempty"` + // LogLevel is the logging level for the session. LogLevel LoggingLevel `json:"logLevel"` diff --git a/mcp/shared.go b/mcp/shared.go index e90bcbd..cacfa75 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -14,8 +14,9 @@ package mcp import ( "context" "encoding/json" + "errors" "fmt" - "log" + "log/slog" "net/http" "reflect" "slices" @@ -23,40 +24,89 @@ import ( "time" "github.com/modelcontextprotocol/go-sdk/auth" + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/internal/mcpgodebug" "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) +// nowrapinvalidparams is a compatibility parameter that restores the previous +// behavior of [methodInfo.unmarshalParams]. When unset (the default), a +// params-decoding failure is wrapped with [jsonrpc2.ErrInvalidParams] so the +// wire response carries error code -32602 ("invalid params") rather than the +// zero-value code 0. See: +// https://github.com/modelcontextprotocol/go-sdk/issues/976#issuecomment-4829124838. +// +// See the documentation for the mcpgodebug package for instructions how to enable it. +// The option will be removed in a future version of the SDK. +var nowrapinvalidparams = mcpgodebug.Value("nowrapinvalidparams") + const ( // latestProtocolVersion is the latest protocol version that this version of // the SDK supports. // // It is the version that the client sends in the initialization request, and // the default version used by the server. - latestProtocolVersion = protocolVersion20250618 + latestProtocolVersion = protocolVersion20260728 + protocolVersion20260728 = "2026-07-28" + protocolVersion20251125 = "2025-11-25" protocolVersion20250618 = "2025-06-18" protocolVersion20250326 = "2025-03-26" protocolVersion20241105 = "2024-11-05" ) var supportedProtocolVersions = []string{ + protocolVersion20260728, + protocolVersion20251125, protocolVersion20250618, protocolVersion20250326, protocolVersion20241105, } +// SupportedProtocolVersions returns the protocol versions supported by this +// version of the SDK, newest first. +// +// Use it to discover the values accepted by +// [ServerOptions.SupportedProtocolVersions]. The returned slice is a copy: +// modifying it does not change what the SDK supports. +func SupportedProtocolVersions() []string { return slices.Clone(supportedProtocolVersions) } + // negotiatedVersion returns the effective protocol version to use, given a -// client version. -func negotiatedVersion(clientVersion string) string { - // In general, prefer to use the clientVersion, but if we don't support the - // client's version, use the latest version. +// client version and the versions the server supports, newest first. +func negotiatedVersion(clientVersion string, supported []string) string { + // In general, prefer to use the clientVersion, but if not supported + // by the server, use the latest version the server supports. // - // This handles the case where a new spec version is released, and the SDK - // does not support it yet. - if !slices.Contains(supportedProtocolVersions, clientVersion) { - return latestProtocolVersion + // Cap the supported versions at the legacy protocolVersion20251125, as this + // method is used by the initialize method which is deprecated in + // version protocolVersion20260728. + if slices.Contains(supported, clientVersion) && clientVersion < protocolVersion20260728 { + return clientVersion } - return clientVersion + for _, v := range supported { + if v < protocolVersion20260728 { + return v + } + } + // If the server was enforced to support only protocolVersion20260728 and + // later, the initialize function should not be called. + // If a client calls the initialize function, return the protocolVersion20251125. + // As according to spec (https://modelcontextprotocol.io/specification/2025-11-25/basic/lifecycle#version-negotiation): + // "If the server supports the requested protocol version, it MUST respond with the same version." + // "Otherwise, the server MUST respond with another protocol version it supports." + // "This SHOULD be the latest version supported by the server." + return protocolVersion20251125 +} + +// negotiateMutuallySupportedVersion returns a protocol version that is supported +// by both the client and the server. +func negotiateMutuallySupportedVersion(supported []string) string { + for _, ver := range supportedProtocolVersions { + if slices.Contains(supported, ver) { + return ver + } + } + return "" } // A MethodHandler handles MCP messages. @@ -86,21 +136,33 @@ func addMiddleware(handlerp *MethodHandler, middleware []Middleware) { } } -func defaultSendingMethodHandler[S Session](ctx context.Context, method string, req Request) (Result, error) { +func defaultSendingMethodHandler(ctx context.Context, method string, req Request) (Result, error) { info, ok := req.GetSession().sendingMethodInfos()[method] if !ok { // This can be called from user code, with an arbitrary value for method. return nil, jsonrpc2.ErrNotHandled } + params := req.GetParams() + if initParams, ok := params.(*InitializeParams); ok { + // Fix the marshaling of initialize params, to work around #607. + // + // The initialize params we produce should never be nil, nor have nil + // capabilities, so any panic here is a bug. + params = initParams.toV2() + } // Notifications don't have results. if strings.HasPrefix(method, "notifications/") { - return nil, req.GetSession().getConn().Notify(ctx, method, req.GetParams()) + return nil, req.GetSession().getConn().Notify(ctx, method, params) } // Create the result to unmarshal into. // The concrete type of the result is the return type of the receiving function. res := info.newResult() - if err := call(ctx, req.GetSession().getConn(), method, req.GetParams(), res); err != nil { - return nil, err + if method == methodSubscriptionsListen { + callSubscriptionsListen(ctx, req.GetSession().getConn(), method, params) + } else { + if err := call(ctx, req.GetSession().getConn(), method, params, res); err != nil { + return nil, err + } } return res, nil } @@ -273,8 +335,18 @@ func newMethodInfo[P paramsPtr[T], R Result, T any](flags methodFlags) methodInf unmarshalParams: func(m json.RawMessage) (Params, error) { var p P if m != nil { - if err := json.Unmarshal(m, &p); err != nil { - return nil, fmt.Errorf("unmarshaling %q into a %T: %w", m, p, err) + if err := internaljson.Unmarshal(m, &p); err != nil { + // Legacy behavior: pre-fix versions surfaced this as a + // plain wrapped error, which caused the wire response to + // carry code 0 instead of -32602. Restore via + // MCPGODEBUG=nowrapinvalidparams=1. + if nowrapinvalidparams == "1" { + return nil, fmt.Errorf("unmarshaling %q into a %T: %w", m, p, err) + } + // Wrap jsonrpc2.ErrInvalidParams so toWireError surfaces + // code -32602 ("invalid params") while preserving the + // descriptive message. + return nil, fmt.Errorf("%w: unmarshaling %q into a %T: %w", jsonrpc2.ErrInvalidParams, m, p, err) } } // We must check missingParamsOK here, in addition to checkRequest, to @@ -329,21 +401,80 @@ func clientSessionMethod[P Params, R Result](f func(*ClientSession, context.Cont } } -// Error codes +// MCP-specific error codes. +const ( + // CodeHeaderMismatch indicates that HTTP headers do not match the corresponding values + // in the request body, or that required headers are missing or malformed. + CodeHeaderMismatch = -32020 + // CodeMissingRequiredClientCapabilities is the JSON-RPC error code defined by + // SEP-2575 for MissingRequiredClientCapabilitiesError. + CodeMissingRequiredClientCapabilities = -32021 + // CodeUnsupportedProtocolVersion is the JSON-RPC error code defined by + // SEP-2575 for UnsupportedProtocolVersionError. + CodeUnsupportedProtocolVersion = -32022 + // CodeURLElicitationRequired indicates that the server requires URL elicitation + // before processing the request. The client should execute the elicitation handler + // with the elicitations provided in the error data. + CodeURLElicitationRequired = -32042 +) + +// CodeResourceNotFound indicates that a requested resource could not be found. +// +// By default, the value is -32602 (Invalid Params), as specified in the +// MCP specification (SEP-2164). To restore the pre-1.7.0 release behavior where the +// error code was -32002, set MCPGODEBUG=customresnotfounderrcode=1. +// +// Deprecated: Use [jsonrpc.CodeInvalidParams] directly. This variable will be +// removed in a future version. +var CodeResourceNotFound int64 = jsonrpc.CodeInvalidParams + +// URLElicitationRequiredError returns an error indicating that URL elicitation is required +// before the request can be processed. The elicitations parameter should contain the +// elicitation requests that must be completed. +func URLElicitationRequiredError(elicitations []*ElicitParams) error { + // Validate that all elicitations are URL mode + for _, elicit := range elicitations { + mode := elicit.Mode + if mode == "" { + mode = "form" // default mode + } + if mode != "url" { + panic(fmt.Sprintf("URLElicitationRequiredError requires all elicitations to be URL mode, got %q", mode)) + } + } + + data, err := json.Marshal(map[string]any{ + "elicitations": elicitations, + }) + if err != nil { + // This should never happen with valid ElicitParams + panic(fmt.Sprintf("failed to marshal elicitations: %v", err)) + } + return &jsonrpc.Error{ + Code: CodeURLElicitationRequired, + Message: "URL elicitation required", + Data: json.RawMessage(data), + } +} + +// Internal error codes const ( - codeResourceNotFound = -32002 // The error code if the method exists and was called properly, but the peer does not support it. + // + // TODO(rfindley): this code is wrong, and we should fix it to be + // consistent with other SDKs. codeUnsupportedMethod = -31001 - // The error code for invalid parameters - codeInvalidParams = -32602 ) // notifySessions calls Notify on all the sessions. // Should be called on a copy of the peer sessions. -func notifySessions[S Session, P Params](sessions []S, method string, params P) { +// The logger must be non-nil. +func notifySessions[S Session, P Params](sessions []S, method string, params P, logger *slog.Logger) { if sessions == nil { return } + // Notify with the background context, so the messages are sent on the + // standalone stream. // TODO: make this timeout configurable, or call handleNotify asynchronously. ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() @@ -353,8 +484,7 @@ func notifySessions[S Session, P Params](sessions []S, method string, params P) for _, s := range sessions { req := newRequest(s, params) if err := handleNotify(ctx, method, req); err != nil { - // TODO(jba): surface this error better - log.Printf("calling %s: %v", method, err) + logger.Warn(fmt.Sprintf("calling %s: %v", method, err)) } } } @@ -395,10 +525,107 @@ func setProgressToken(p Params, pt any) { m := p.GetMeta() if m == nil { m = map[string]any{} + p.SetMeta(m) } m[progressTokenKey] = pt } +// extractRequestMeta performs a lightweight partial unmarshal of the `_meta` +// field from a JSON-RPC request's raw params. +func extractRequestMeta(rawParams json.RawMessage) Meta { + if len(rawParams) == 0 { + return nil + } + var meta struct { + Meta Meta `json:"_meta"` + } + if err := internaljson.Unmarshal(rawParams, &meta); err != nil { + return nil + } + return meta.Meta +} + +type validatedMeta struct { + usesNewProtocol bool + initializeParams *InitializeParams + logLevel LoggingLevel +} + +// validateRequestMeta inspects a JSON-RPC request to detect whether it follows +// the >= 2026-07-28 protocol via the `_meta` field. +// If the request has no _meta, no protocolVersion in _meta, or a protocolVersion +// naming a supported version that predates 2026-07-28, it returns a non-nil +// validatedMeta with usesNewProtocol set to false, and a nil error: such a request +// belongs to the initialize handshake. +// If the protocolVersion names a version that is not supported at all, it returns +// usesNewProtocol set to true with initializeParams carrying only that version, so +// that the caller answers with [CodeUnsupportedProtocolVersion]. The rest of the +// `_meta` triple is not validated in that case, since the revision defining it is +// unknown. +// Otherwise, it validates the presence of clientCapabilities in _meta. If it is +// missing or invalid, it returns nil and a non-nil error. clientInfo is optional; +// if present but invalid, an error is returned. Otherwise, it returns +// usesNewProtocol set to true and the populated initializeParams. +// Notifications always report usesNewProtocol false and never error. +func validateRequestMeta(req *jsonrpc.Request) (*validatedMeta, error) { + // SEP-2575 defines the `_meta` triple for calls only: NotificationParams + // declares `_meta` optional with no protocolVersion, so a notification + // neither selects the new protocol nor is discarded for carrying an + // incomplete triple. + if !req.IsCall() { + return &validatedMeta{usesNewProtocol: false, initializeParams: nil}, nil + } + meta := extractRequestMeta(req.Params) + if meta == nil { + return &validatedMeta{usesNewProtocol: false, initializeParams: nil}, nil + } + protocolVersion, ok := meta[MetaKeyProtocolVersion].(string) + // An absent or empty value declares no version at all, which is also how the + // streamable handler reads `_meta.protocolVersion`. + if !ok || protocolVersion == "" { + return &validatedMeta{usesNewProtocol: false, initializeParams: nil}, nil + } + // Classify by membership rather than by ordering: a version the SDK does not + // know sorts arbitrarily against 2026-07-28, and one that sorts below it would + // be served as a legacy handshake instead of reaching the + // CodeUnsupportedProtocolVersion check in ServerSession.handle. The + // streamable handler already applies the same membership rule to the + // Mcp-Protocol-Version header. + supported := slices.Contains(supportedProtocolVersions, protocolVersion) + if supported && protocolVersion < protocolVersion20260728 { + return &validatedMeta{usesNewProtocol: false, initializeParams: nil}, nil + } + if !supported { + return &validatedMeta{usesNewProtocol: true, initializeParams: &InitializeParams{ + ProtocolVersion: protocolVersion, + }}, nil + } + var clientInfo *Implementation + if _, present := meta[MetaKeyClientInfo]; present { + var ok bool + clientInfo, ok = decodeMetaValue[*Implementation](meta, MetaKeyClientInfo) + if !ok { + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: fmt.Sprintf("invalid _meta field %q", MetaKeyClientInfo), + } + } + } + capabilities, ok := decodeMetaValue[*clientCapabilitiesV2](meta, MetaKeyClientCapabilities) + if !ok { + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: fmt.Sprintf("missing or invalid _meta field %q", MetaKeyClientCapabilities), + } + } + logLevel, _ := decodeMetaValue[LoggingLevel](meta, MetaKeyLogLevel) + return &validatedMeta{usesNewProtocol: true, initializeParams: &InitializeParams{ + ProtocolVersion: protocolVersion, + Capabilities: capabilities.toV1(), + ClientInfo: clientInfo, + }, logLevel: logLevel}, nil +} + // A Request is a method request with parameters and additional information, such as the session. // Request is implemented by [*ClientRequest] and [*ServerRequest]. type Request interface { @@ -427,6 +654,26 @@ type ServerRequest[P Params] struct { type RequestExtra struct { TokenInfo *auth.TokenInfo // bearer token info (e.g. from OAuth) if any Header http.Header // header from HTTP request, if any + + // If set, CloseSSEStream explicitly closes the current SSE request stream. + // + // [SEP-1699] introduced server-side SSE stream disconnection: for + // long-running requests, servers may opt to close the SSE stream and + // ask the client to retry at a later time. CloseSSEStream implements this + // feature; if RetryAfter is set, an event is sent with a `retry:` field + // to configure the reconnection delay. + // + // [SEP-1699]: https://github.com/modelcontextprotocol/modelcontextprotocol/issues/1699 + // This mechanism is deprecated in protocol version 2026-07-28 as the resumability + // feature is removed. + CloseSSEStream func(CloseSSEStreamArgs) +} + +// CloseSSEStreamArgs are arguments for [RequestExtra.CloseSSEStream]. +type CloseSSEStreamArgs struct { + // RetryAfter configures the reconnection delay sent to the client via the + // SSE retry field. If zero, no retry field is sent. + RetryAfter time.Duration } func (*ClientRequest[P]) isRequest() {} @@ -441,6 +688,117 @@ func (r *ServerRequest[P]) GetParams() Params { return r.Params } func (r *ClientRequest[P]) GetExtra() *RequestExtra { return nil } func (r *ServerRequest[P]) GetExtra() *RequestExtra { return r.Extra } +// HasParams reports whether req carries params. +// +// The "params" member of a message may be omitted for several methods, among +// them the list methods, ping and notifications/initialized. For such a request +// [Request.GetParams] returns a non-nil [Params] holding a nil pointer, so +// comparing its result against nil reports that params are present, while +// calling a method on that result panics. Middleware that inspects params must +// ask this first: +// +// if mcp.HasParams(req) { +// log.Println(method, req.GetParams().GetMeta()) +// } +func HasParams(req Request) bool { + params := req.GetParams() + if params == nil { + return false + } + // isNil cannot answer this for every value: a params type embedding + // ParamsBase promotes isNil through a field selector, which dereferences + // the nil outer pointer before the method body runs. + v := reflect.ValueOf(params) + return v.Kind() != reflect.Pointer || !v.IsNil() +} + +// ProtocolVersion returns the protocol version negotiated for this request. +// +// For requests following the >= 2026-07-28 protocol, the value is read from +// the per-request `_meta` field. For older protocol requests, the value falls +// back to the session-level [InitializeParams] established during the +// initialize handshake. +func (r *ServerRequest[P]) ProtocolVersion() string { + if m := getRequestMeta(r); m != nil { + if v, ok := m[MetaKeyProtocolVersion].(string); ok { + return v + } + } + if r.Session != nil { + if p := r.Session.InitializeParams(); p != nil { + return p.ProtocolVersion + } + } + return "" +} + +// ClientInfo returns the [Implementation] identifying the calling client. +// +// For requests following the >= 2026-07-28 protocol, the value is read from +// the per-request `_meta` field. For older protocol requests, the value falls +// back to the session-level [InitializeParams]. +func (r *ServerRequest[P]) ClientInfo() *Implementation { + if m := getRequestMeta(r); m != nil { + if v, ok := decodeMetaValue[*Implementation](m, MetaKeyClientInfo); ok { + return v + } + } + if r.Session != nil { + if p := r.Session.InitializeParams(); p != nil { + return p.ClientInfo + } + } + return nil +} + +// ClientCapabilities returns the [ClientCapabilities] of the calling client. +// +// For requests following the >= 2026-07-28 protocol, the value is read from +// the per-request `_meta` field. For older protocol requests, the value falls +// back to the session-level [InitializeParams]. +func (r *ServerRequest[P]) ClientCapabilities() *ClientCapabilities { + if m := getRequestMeta(r); m != nil { + if v, ok := decodeMetaValue[*clientCapabilitiesV2](m, MetaKeyClientCapabilities); ok { + return v.toV1() + } + } + if r.Session != nil { + if p := r.Session.InitializeParams(); p != nil { + return p.Capabilities + } + } + return nil +} + +// getRequestMeta returns the raw `_meta` map from the request's params, or +// nil if the params are absent. +func getRequestMeta[P Params](r *ServerRequest[P]) map[string]any { + if !HasParams(r) { + return nil + } + return r.Params.GetMeta() +} + +// decodeMetaValue decodes a typed value out of a `_meta` map. Values may +// arrive either as the typed Go value (when constructed in-process) or as +// the generic JSON map produced by encoding/json after wire transit. In the +// latter case, the value is re-encoded and decoded into the target type. +func decodeMetaValue[T any](m map[string]any, key string) (T, bool) { + var zero T + raw, ok := m[key] + if !ok || raw == nil { + return zero, false + } + if v, ok := raw.(T); ok { + return v, true + } + var v T + if err := remarshal(raw, &v); err != nil { + return zero, false + } + return v, true +} + func serverRequestFor[P Params](s *ServerSession, p P) *ServerRequest[P] { return &ServerRequest[P]{Session: s, Params: p} } @@ -458,8 +816,27 @@ type Params interface { // isParams discourages implementation of Params outside of this package. isParams() + + // isNil returns true if the underlying value is nil. + isNil() bool } +// ParamsBase can be embedded in custom parameter structs to satisfy the +// [Params] interface. It provides the required [Meta] field and the unexported +// isParams marker method. +// +// type SearchParams struct { +// mcp.ParamsBase +// Query string `json:"query"` +// } +type ParamsBase struct { + Meta `json:"_meta,omitempty"` +} + +func (*ParamsBase) isParams() {} + +func (p *ParamsBase) isNil() bool { return p == nil } + // RequestParams is a parameter (input) type for an MCP request. type RequestParams interface { Params @@ -478,17 +855,36 @@ type Result interface { // isResult discourages implementation of Result outside of this package. isResult() + // isNil returns true if the underlying value is nil. + isNil() bool + // GetMeta returns metadata from a value. GetMeta() map[string]any // SetMeta sets the metadata on a value. SetMeta(map[string]any) } +// ResultBase can be embedded in custom result structs to satisfy the +// [Result] interface. It provides the required [Meta] field and the unexported +// isResult marker method. +// +// type SearchResult struct { +// mcp.ResultBase +// Hits []string `json:"hits"` +// } +type ResultBase struct { + Meta `json:"_meta,omitempty"` +} + +func (*ResultBase) isResult() {} +func (x *ResultBase) isNil() bool { return x == nil } + // emptyResult is returned by methods that have no result, like ping. // Those methods cannot return nil, because jsonrpc2 cannot handle nils. type emptyResult struct{} func (*emptyResult) isResult() {} +func (x *emptyResult) isNil() bool { return x == nil } func (*emptyResult) GetMeta() map[string]any { panic("should never be called") } func (*emptyResult) SetMeta(map[string]any) { panic("should never be called") } @@ -511,7 +907,21 @@ type keepaliveSession interface { // startKeepalive starts the keepalive mechanism for a session. // It assigns the cancel function to the provided cancelPtr and starts a goroutine // that sends ping messages at the specified interval. -func startKeepalive(session keepaliveSession, interval time.Duration, cancelPtr *context.CancelFunc) { +// +// failureThreshold is the number of consecutive ping failures tolerated before +// the session is closed; a value below 1 is treated as 1 (close on the first +// failure). A successful ping resets the counter. This mirrors the spec's +// "multiple failed pings MAY trigger a connection reset" language, letting a +// transient miss pass without tearing down an otherwise live session. +// +// logger must be non-nil; ping failures (both the tolerated ones and the final +// one that closes the session) are reported via logger so they are not silently +// dropped. +func startKeepalive(session keepaliveSession, interval time.Duration, failureThreshold int, cancelPtr *context.CancelFunc, logger *slog.Logger) { + if failureThreshold < 1 { + failureThreshold = 1 + } + ctx, cancel := context.WithCancel(context.Background()) // Assign cancel function before starting goroutine to avoid race condition. // We cannot return it because the caller may need to cancel during the @@ -522,6 +932,7 @@ func startKeepalive(session keepaliveSession, interval time.Duration, cancelPtr ticker := time.NewTicker(interval) defer ticker.Stop() + consecutiveFailures := 0 for { select { case <-ctx.Done(): @@ -530,11 +941,32 @@ func startKeepalive(session keepaliveSession, interval time.Duration, cancelPtr pingCtx, pingCancel := context.WithTimeout(context.Background(), interval/2) err := session.Ping(pingCtx, nil) pingCancel() - if err != nil { - // Ping failed, close the session - _ = session.Close() + if err == nil { + consecutiveFailures = 0 + continue + } + if errors.Is(err, jsonrpc2.ErrMethodNotFound) { + // Peer doesn't support ping, stop the keepalive process. return } + consecutiveFailures++ + if consecutiveFailures < failureThreshold { + // Tolerate transient failures below the threshold; log so + // the misses are still observable to operators. See #218. + logger.Warn("keepalive ping failed; tolerating below threshold", + "error", err, + "consecutiveFailures", consecutiveFailures, + "failureThreshold", failureThreshold) + continue + } + // Threshold reached; log before closing the session so the + // failure is observable to operators. See #218. + logger.Error("keepalive ping failed; closing session", + "error", err, + "consecutiveFailures", consecutiveFailures, + "failureThreshold", failureThreshold) + _ = session.Close() + return } } }() diff --git a/mcp/shared_test.go b/mcp/shared_test.go index 23818f8..4566768 100644 --- a/mcp/shared_test.go +++ b/mcp/shared_test.go @@ -4,6 +4,415 @@ package mcp +import ( + "encoding/json" + "errors" + "strings" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" +) + +func TestValidateRequestMeta(t *testing.T) { + tests := []struct { + name string + method string + params any + wantUsesNew bool + wantLogLevel LoggingLevel + wantErrContains string + }{ + { + name: "no params: old protocol", + method: methodListTools, + params: nil, + wantUsesNew: false, + }, + { + name: "no _meta: old protocol", + method: methodCallTool, + params: map[string]any{"name": "x"}, + wantUsesNew: false, + }, + { + name: "_meta without protocolVersion: old protocol", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{"otherKey": "v"}, + "name": "x", + }, + wantUsesNew: false, + }, + { + name: "new protocol with all required fields", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "c", "version": "1"}, + MetaKeyClientCapabilities: map[string]any{}, + }, + "name": "x", + }, + wantUsesNew: true, + }, + { + name: "new protocol missing clientInfo", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientCapabilities: map[string]any{}, + }, + "name": "x", + }, + wantUsesNew: true, + }, + { + name: "new protocol invalid clientInfo", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: "not an object", + MetaKeyClientCapabilities: map[string]any{}, + }, + "name": "x", + }, + wantUsesNew: false, + wantErrContains: MetaKeyClientInfo, + }, + { + name: "new protocol missing clientCapabilities", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "c", "version": "1"}, + }, + "name": "x", + }, + wantUsesNew: false, + wantErrContains: MetaKeyClientCapabilities, + }, + { + name: "malformed _meta is ignored", + method: methodCallTool, + params: json.RawMessage(`{"_meta": "not an object", "name": "x"}`), + wantUsesNew: false, + }, + { + name: "supported pre-2026-07-28 version: old protocol", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20251125, + }, + "name": "x", + }, + wantUsesNew: false, + }, + { + name: "empty protocolVersion declares nothing: old protocol", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: "", + }, + "name": "x", + }, + wantUsesNew: false, + }, + { + // An unsupported version sorting below 2026-07-28 is still a + // new-protocol request, so that the caller can answer it with + // CodeUnsupportedProtocolVersion. + name: "unsupported version: new protocol", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: "1900-01-01", + MetaKeyClientInfo: map[string]any{"name": "c", "version": "1"}, + MetaKeyClientCapabilities: map[string]any{}, + }, + "name": "x", + }, + wantUsesNew: true, + }, + { + // The triple is defined by a revision the SDK does not know, so an + // unsupported version is reported without validating the rest of it. + name: "unsupported version missing clientCapabilities: new protocol", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: "1900-01-01", + }, + "name": "x", + }, + wantUsesNew: true, + }, + { + name: "new protocol with logLevel", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "c", "version": "1"}, + MetaKeyClientCapabilities: map[string]any{}, + MetaKeyLogLevel: "warning", + }, + "name": "x", + }, + wantUsesNew: true, + wantLogLevel: "warning", + }, + { + name: "new protocol without logLevel", + method: methodCallTool, + params: map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "c", "version": "1"}, + MetaKeyClientCapabilities: map[string]any{}, + }, + "name": "x", + }, + wantUsesNew: true, + wantLogLevel: "", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var raw json.RawMessage + switch p := tc.params.(type) { + case json.RawMessage: + raw = p + default: + raw = mustMarshal(tc.params) + } + id, err := jsonrpc.MakeID("test") + if err != nil { + t.Fatal(err) + } + req := &jsonrpc.Request{Method: tc.method, Params: raw, ID: id} + + vmeta, err := validateRequestMeta(req) + usesNew := vmeta != nil && vmeta.usesNewProtocol + if usesNew != tc.wantUsesNew { + t.Errorf("usesNewProtocol = %v, want %v", usesNew, tc.wantUsesNew) + } + if vmeta != nil && vmeta.logLevel != tc.wantLogLevel { + t.Errorf("logLevel = %q, want %q", vmeta.logLevel, tc.wantLogLevel) + } + if tc.wantErrContains == "" { + if err != nil { + t.Errorf("unexpected error: %v", err) + } + return + } + if err == nil { + t.Fatalf("expected error containing %q, got nil", tc.wantErrContains) + } + var jerr *jsonrpc.Error + if !errors.As(err, &jerr) { + t.Fatalf("expected *jsonrpc.Error, got %T: %v", err, err) + } + if jerr.Code != jsonrpc.CodeInvalidParams { + t.Errorf("error code = %d, want %d", jerr.Code, jsonrpc.CodeInvalidParams) + } + if !strings.Contains(jerr.Message, tc.wantErrContains) { + t.Errorf("error message %q does not contain %q", jerr.Message, tc.wantErrContains) + } + }) + } +} + +func TestServerRequest_PerRequestAccessors(t *testing.T) { + // A request carrying the new-protocol _meta fields populates the + // accessors with values from _meta. + caps := &ClientCapabilities{Sampling: &SamplingCapabilities{}} + info := &Implementation{Name: "c", Version: "1"} + params := &CallToolParamsRaw{ + Meta: Meta{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: info, + MetaKeyClientCapabilities: caps, + }, + Name: "x", + } + req := &ServerRequest[*CallToolParamsRaw]{Params: params} + if got := req.ProtocolVersion(); got != protocolVersion20260728 { + t.Errorf("ProtocolVersion = %q, want %q", got, protocolVersion20260728) + } + if got := req.ClientInfo(); got == nil || got.Name != "c" { + t.Errorf("ClientInfo = %+v, want Name=c", got) + } + if got := req.ClientCapabilities(); got == nil || got.Sampling == nil { + t.Errorf("ClientCapabilities = %+v, want non-nil Sampling", got) + } +} + +func TestServerRequest_PerRequestAccessors_FromJSON(t *testing.T) { + // Values arriving over the wire are JSON maps; the accessors should + // re-decode them into typed Go values. + raw := json.RawMessage(`{ + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": {"name": "wire-client", "version": "9"}, + "io.modelcontextprotocol/clientCapabilities": {"sampling": {}} + }, + "name": "tool" + }`) + var params CallToolParamsRaw + if err := json.Unmarshal(raw, ¶ms); err != nil { + t.Fatal(err) + } + req := &ServerRequest[*CallToolParamsRaw]{Params: ¶ms} + if got, want := req.ProtocolVersion(), protocolVersion20260728; got != want { + t.Errorf("ProtocolVersion = %q, want %q", got, want) + } + gotInfo := req.ClientInfo() + wantInfo := &Implementation{Name: "wire-client", Version: "9"} + if diff := cmp.Diff(wantInfo, gotInfo); diff != "" { + t.Errorf("ClientInfo mismatch (-want +got):\n%s", diff) + } + gotCaps := req.ClientCapabilities() + if gotCaps == nil || gotCaps.Sampling == nil { + t.Errorf("ClientCapabilities = %+v, want non-nil Sampling", gotCaps) + } +} + +func TestServerRequest_PerRequestAccessors_FallbackToInitializeParams(t *testing.T) { + // With no _meta on the request, accessors must fall back to the + // session's InitializeParams (the old-protocol path). + ss := &ServerSession{} + ss.state.InitializeParams = &InitializeParams{ + ProtocolVersion: protocolVersion20251125, + ClientInfo: &Implementation{Name: "old", Version: "0"}, + Capabilities: &ClientCapabilities{Elicitation: &ElicitationCapabilities{}}, + } + req := &ServerRequest[*CallToolParamsRaw]{ + Session: ss, + Params: &CallToolParamsRaw{Name: "x"}, + } + if got, want := req.ProtocolVersion(), protocolVersion20251125; got != want { + t.Errorf("ProtocolVersion fallback = %q, want %q", got, want) + } + if got := req.ClientInfo(); got == nil || got.Name != "old" { + t.Errorf("ClientInfo fallback = %+v, want Name=old", got) + } + if got := req.ClientCapabilities(); got == nil || got.Elicitation == nil { + t.Errorf("ClientCapabilities fallback = %+v, want non-nil Elicitation", got) + } +} + +func TestServerRequest_PerRequestAccessors_Empty(t *testing.T) { + // With no _meta and no session, accessors return zero values. + req := &ServerRequest[*CallToolParamsRaw]{ + Params: &CallToolParamsRaw{Name: "x"}, + } + if got := req.ProtocolVersion(); got != "" { + t.Errorf("ProtocolVersion = %q, want empty", got) + } + if got := req.ClientInfo(); got != nil { + t.Errorf("ClientInfo = %+v, want nil", got) + } + if got := req.ClientCapabilities(); got != nil { + t.Errorf("ClientCapabilities = %+v, want nil", got) + } +} + +func TestHasParams(t *testing.T) { + // {"jsonrpc":"2.0","id":1,"method":"tools/list"} is a complete request: + // serverMethodInfos marks these methods missingParamsOK. Build what + // handleReceive hands the middleware chain for one. + for _, method := range []string{ + methodListTools, + methodListPrompts, + methodListResources, + notificationInitialized, + } { + t.Run(method, func(t *testing.T) { + info := serverMethodInfos[method] + params, err := info.unmarshalParams(nil) + if err != nil { + t.Fatal(err) + } + req := info.newRequest(&ServerSession{}, params, nil) + if req.GetParams() == nil { + t.Fatal("GetParams returned a nil interface, so this test no longer covers the typed-nil case") + } + if HasParams(req) { + t.Error("HasParams = true for a request whose params member was omitted") + } + }) + } + + t.Run("params present", func(t *testing.T) { + info := serverMethodInfos[methodListTools] + params, err := info.unmarshalParams(json.RawMessage(`{"cursor":"c"}`)) + if err != nil { + t.Fatal(err) + } + req := info.newRequest(&ServerSession{}, params, nil) + if !HasParams(req) { + t.Error("HasParams = false for a request carrying params") + } + }) +} + +func TestServerRequest_PerRequestAccessors_MissingParams(t *testing.T) { + // Custom methods also allow the params member to be omitted, and their + // params types embed ParamsBase, whose promoted isNil dereferences the + // typed-nil outer pointer. + type customParams struct{ ParamsBase } + req := &ServerRequest[*customParams]{} + if got := req.ProtocolVersion(); got != "" { + t.Errorf("ProtocolVersion = %q, want empty", got) + } + if got := req.ClientInfo(); got != nil { + t.Errorf("ClientInfo = %+v, want nil", got) + } + if got := req.ClientCapabilities(); got != nil { + t.Errorf("ClientCapabilities = %+v, want nil", got) + } +} + +func TestImplementationDescriptionJSON(t *testing.T) { + impl := &Implementation{ + Name: "greeter", + Title: "Greeter", + Description: "Example server for greeting tools", + Version: "v1.0.0", + } + got, err := json.Marshal(impl) + if err != nil { + t.Fatal(err) + } + want := `{"name":"greeter","title":"Greeter","description":"Example server for greeting tools","version":"v1.0.0"}` + if string(got) != want { + t.Fatalf("Implementation JSON = %s, want %s", got, want) + } + + var roundTrip Implementation + if err := json.Unmarshal(got, &roundTrip); err != nil { + t.Fatal(err) + } + if diff := cmp.Diff(impl, &roundTrip); diff != "" { + t.Fatalf("Implementation round trip mismatch (-want +got):\n%s", diff) + } + + got, err = json.Marshal(&Implementation{Name: "greeter", Version: "v1.0.0"}) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(got), "description") { + t.Fatalf("empty description should be omitted, got %s", got) + } +} + // TODO(v0.3.0): rewrite this test. // func TestToolValidate(t *testing.T) { // // Check that the tool returned from NewServerTool properly validates its input schema. @@ -223,3 +632,29 @@ package mcp // } // }) // } + +// TestUnmarshalParamsInvalidParamsCode is a regression test for +// https://github.com/modelcontextprotocol/go-sdk/issues/976#issuecomment-4829124838. +// Malformed params for a registered method must surface as a JSON-RPC error +// whose code is CodeInvalidParams (-32602) rather than the zero-value 0. +func TestUnmarshalParamsInvalidParamsCode(t *testing.T) { + info, ok := serverMethodInfos[methodPing] + if !ok { + t.Fatalf("no methodInfo for %q", methodPing) + } + // Array where a struct is expected. + _, err := info.unmarshalParams(json.RawMessage(`["a","b"]`)) + if err == nil { + t.Fatal("unmarshalParams returned nil error for malformed params") + } + var jerr *jsonrpc.Error + if !errors.As(err, &jerr) { + t.Fatalf("expected *jsonrpc.Error in chain, got %T: %v", err, err) + } + if jerr.Code != jsonrpc.CodeInvalidParams { + t.Errorf("error code = %d, want %d (CodeInvalidParams)", jerr.Code, jsonrpc.CodeInvalidParams) + } + if !strings.Contains(err.Error(), "unmarshaling") { + t.Errorf("error message = %q, want it to mention unmarshaling", err.Error()) + } +} diff --git a/mcp/sse.go b/mcp/sse.go index 7f64491..f5e90b8 100644 --- a/mcp/sse.go +++ b/mcp/sse.go @@ -7,13 +7,18 @@ package mcp import ( "bytes" "context" + "crypto/rand" + "errors" "fmt" "io" + "mime" + "net" "net/http" "net/url" "sync" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/internal/util" "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) @@ -51,9 +56,23 @@ type SSEHandler struct { } // SSEOptions specifies options for an [SSEHandler]. -// for now, it is empty, but may be extended in future. -// https://github.com/modelcontextprotocol/go-sdk/issues/507 -type SSEOptions struct{} +type SSEOptions struct { + // DisableLocalhostProtection disables automatic DNS rebinding protection. + // By default, requests arriving via a localhost address (127.0.0.1, [::1]) + // that have a non-localhost Host header are rejected with 403 Forbidden. + // This protects against DNS rebinding attacks regardless of whether the + // server is listening on localhost specifically or on 0.0.0.0. + // + // Only disable this if you understand the security implications. + // See: https://modelcontextprotocol.io/specification/2025-11-25/basic/security_best_practices#local-mcp-server-compromise + DisableLocalhostProtection bool + + // MaxRequestBodyBytes limits the number of bytes read from an incoming + // POST body before the server responds with 413 Request Entity Too Large. + // It matches [StreamableHTTPOptions.MaxRequestBodyBytes]: if zero, + // [DefaultMaxRequestBodyBytes] is used; a negative value disables the limit. + MaxRequestBodyBytes int64 +} // NewSSEHandler returns a new [SSEHandler] that creates and manages MCP // sessions created via incoming HTTP requests. @@ -110,6 +129,12 @@ type SSEServerTransport struct { // Response is the hanging response body to the incoming GET request. Response http.ResponseWriter + // MaxRequestBodyBytes limits the number of bytes read from a POSTed + // message body before [SSEServerTransport.ServeHTTP] responds with 413 + // Request Entity Too Large. If zero, [DefaultMaxRequestBodyBytes] is used; + // a negative value disables the limit. See [SSEOptions.MaxRequestBodyBytes]. + MaxRequestBodyBytes int64 + // incoming is the queue of incoming messages. // It is never closed, and by convention, incoming is non-nil if and only if // the transport is connected. @@ -131,9 +156,22 @@ func (t *SSEServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request) return } + limit := t.MaxRequestBodyBytes + if limit == 0 { + limit = DefaultMaxRequestBodyBytes + } + if limit > 0 && req.Body != nil { + req.Body = http.MaxBytesReader(w, req.Body, limit) + } + // Read and parse the message. data, err := io.ReadAll(req.Body) if err != nil { + var mbe *http.MaxBytesError + if errors.As(err, &mbe) { + http.Error(w, fmt.Sprintf("request body exceeds %d bytes", mbe.Limit), http.StatusRequestEntityTooLarge) + return + } http.Error(w, "failed to read body", http.StatusBadRequest) return } @@ -177,10 +215,35 @@ func (t *SSEServerTransport) Connect(context.Context) (Connection, error) { return &sseServerConn{t: t}, nil } -func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { - sessionID := req.URL.Query().Get("sessionid") +// SupportsProtocolVersion reports whether the HTTP+SSE transport can serve the +// given protocol version. MCP 2026-07-28 defines only the stdio and Streamable +// HTTP bindings, so this (deprecated) transport does not advertise that revision. +func (t *SSEServerTransport) SupportsProtocolVersion(version string) bool { + return version < protocolVersion20260728 +} - // TODO: consider checking Content-Type here. For now, we are lax. +func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { + // DNS rebinding protection: auto-enabled for localhost servers. + // See: https://modelcontextprotocol.io/specification/2025-11-25/basic/security_best_practices#local-mcp-server-compromise + if !h.opts.DisableLocalhostProtection { + if localAddr, ok := req.Context().Value(http.LocalAddrContextKey).(net.Addr); ok && localAddr != nil { + if util.IsLoopback(localAddr.String()) && !util.IsLoopback(req.Host) { + http.Error(w, fmt.Sprintf("Forbidden: invalid Host header %q", req.Host), http.StatusForbidden) + return + } + } + } + + // Validate 'Content-Type' header. + if req.Method == http.MethodPost { + mediaType, _, err := mime.ParseMediaType(req.Header.Get("Content-Type")) + if err != nil || mediaType != "application/json" { + http.Error(w, "Content-Type must be 'application/json'", http.StatusUnsupportedMediaType) + return + } + } + + sessionID := req.URL.Query().Get("sessionid") // For POST requests, the message body is a message to send to a session. if req.Method == http.MethodPost { @@ -202,7 +265,8 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { } if req.Method != http.MethodGet { - http.Error(w, "invalid method", http.StatusMethodNotAllowed) + w.Header().Set("Allow", "GET, POST") + http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed) return } @@ -215,14 +279,18 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") - sessionID = randText() + sessionID = rand.Text() endpoint, err := req.URL.Parse("?sessionid=" + sessionID) if err != nil { http.Error(w, "internal error: failed to create endpoint", http.StatusInternalServerError) return } - transport := &SSEServerTransport{Endpoint: endpoint.RequestURI(), Response: w} + transport := &SSEServerTransport{ + Endpoint: endpoint.RequestURI(), + Response: w, + MaxRequestBodyBytes: h.opts.MaxRequestBodyBytes, + } // The session is terminated when the request exits. h.mu.Lock() @@ -329,6 +397,11 @@ type SSEClientTransport struct { // HTTPClient is the client to use for making HTTP requests. If nil, // http.DefaultClient is used. HTTPClient *http.Client + + // MaxEventSize bounds the number of bytes buffered while reading a single + // server-sent event. A value of 0 selects [DefaultMaxEventSize], a negative + // value disables the cap. + MaxEventSize int } // Connect connects through the client endpoint. @@ -351,9 +424,24 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) { return nil, err } + // Check HTTP status code before attempting to parse SSE events. + // This ensures proper error reporting for authentication failures (401), + // authorization failures (403), and other HTTP errors. + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + resp.Body.Close() + return nil, fmt.Errorf("failed to connect: %s", http.StatusText(resp.StatusCode)) + } + + maxEventSize := c.MaxEventSize + if maxEventSize == 0 { + maxEventSize = DefaultMaxEventSize + } + + // Reuse the scanner so events buffered after the endpoint are not lost. + events := scanEventsLimited(resp.Body, maxEventSize) msgEndpoint, err := func() (*url.URL, error) { var evt Event - for evt, err = range scanEvents(resp.Body) { + for evt, err = range events { break } if err != nil { @@ -382,7 +470,7 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) { go func() { defer s.Close() // close the transport when the GET exits - for evt, err := range scanEvents(resp.Body) { + for evt, err := range events { if err != nil { return } diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index 6132d31..434f6de 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -27,6 +27,8 @@ func Add(ctx context.Context, req *mcp.CallToolRequest, args AddParams) (*mcp.Ca }, nil, nil } +// !+ssehandler + func ExampleSSEHandler() { server := mcp.NewServer(&mcp.Implementation{Name: "adder", Version: "v0.0.1"}, nil) mcp.AddTool(server, &mcp.Tool{Name: "add", Description: "add two numbers"}, Add) @@ -55,3 +57,5 @@ func ExampleSSEHandler() { // Output: 3 } + +// !-ssehandler diff --git a/mcp/sse_test.go b/mcp/sse_test.go index 25435ff..d9b7a2b 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -9,14 +9,82 @@ import ( "context" "fmt" "io" + "net" "net/http" "net/http/httptest" + "slices" + "strings" "sync/atomic" "testing" + "time" "github.com/google/go-cmp/cmp" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) +// TestSSEServerTransport_SupportedVersions verifies that the deprecated +// HTTP+SSE transport does not advertise protocol revisions that only define +// stdio and Streamable HTTP (see #1112). +func TestSSEServerTransport_SupportedVersions(t *testing.T) { + var tr SSEServerTransport + versions := filterSupportedVersions(&tr, supportedProtocolVersions) + if slices.Contains(versions, protocolVersion20260728) { + t.Errorf("filterSupportedVersions(SSEServerTransport) = %v, must not include %q", + versions, protocolVersion20260728) + } + if !slices.Contains(versions, protocolVersion20241105) { + t.Errorf("filterSupportedVersions(SSEServerTransport) = %v, want %q (SSE-defined revision)", + versions, protocolVersion20241105) + } + if len(versions) == 0 { + t.Error("filterSupportedVersions(SSEServerTransport) is empty; want legacy SSE revisions") + } + for _, v := range supportedProtocolVersions { + got := tr.SupportsProtocolVersion(v) + want := v < protocolVersion20260728 + if got != want { + t.Errorf("SupportsProtocolVersion(%q) = %v, want %v", v, got, want) + } + } +} + +// TestSSEHandler_DiscoverFallsBackToInitialize checks that a Modern-first +// client connecting over HTTP+SSE does not stick on 2026-07-28: discover +// advertises only SSE-defined revisions, and Connect falls back to legacy +// initialize successfully. +func TestSSEHandler_DiscoverFallsBackToInitialize(t *testing.T) { + ctx := context.Background() + server := NewServer(testImpl, nil) + sseHandler := NewSSEHandler(func(*http.Request) *Server { return server }, nil) + httpServer := httptest.NewServer(sseHandler) + defer httpServer.Close() + + client := NewClient(testImpl, nil) + // Default Client.Connect probes with the latest protocol version (2026-07-28). + cs, err := client.Connect(ctx, &SSEClientTransport{Endpoint: httpServer.URL}, nil) + if err != nil { + t.Fatalf("Connect over SSE with modern client: %v", err) + } + defer cs.Close() + + ir := cs.InitializeResult() + if ir == nil { + t.Fatal("InitializeResult is nil after Connect") + } + if ir.ProtocolVersion == protocolVersion20260728 { + t.Errorf("InitializeResult.ProtocolVersion = %q; HTTP+SSE must not negotiate that revision", + ir.ProtocolVersion) + } + if ir.ProtocolVersion != protocolVersion20251125 { + // Client.Connect falls back to protocolVersion20251125 for legacy initialize. + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (legacy fallback)", + ir.ProtocolVersion, protocolVersion20251125) + } + if err := cs.Ping(ctx, nil); err != nil { + t.Errorf("Ping after SSE fallback initialize: %v", err) + } +} + func TestSSEServer(t *testing.T) { for _, closeServerFirst := range []bool{false, true} { t.Run(fmt.Sprintf("closeServerFirst=%t", closeServerFirst), func(t *testing.T) { @@ -131,3 +199,285 @@ type roundTripperFunc func(*http.Request) (*http.Response, error) func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } + +func TestSSEClientTransport_HTTPErrors(t *testing.T) { + tests := []struct { + name string + statusCode int + wantErrContain string + }{ + { + name: "401 Unauthorized", + statusCode: http.StatusUnauthorized, + wantErrContain: "Unauthorized", + }, + { + name: "403 Forbidden", + statusCode: http.StatusForbidden, + wantErrContain: "Forbidden", + }, + { + name: "404 Not Found", + statusCode: http.StatusNotFound, + wantErrContain: "Not Found", + }, + { + name: "500 Internal Server Error", + statusCode: http.StatusInternalServerError, + wantErrContain: "Internal Server Error", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create a test server that returns the specified status code + httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, http.StatusText(tt.statusCode), tt.statusCode) + })) + defer httpServer.Close() + + clientTransport := &SSEClientTransport{ + Endpoint: httpServer.URL, + } + + c := NewClient(testImpl, nil) + _, err := c.Connect(context.Background(), clientTransport, nil) + + if err == nil { + t.Fatalf("expected error, got nil") + } + + errStr := err.Error() + if !bytes.Contains([]byte(errStr), []byte(tt.wantErrContain)) { + t.Errorf("error message %q does not contain %q", errStr, tt.wantErrContain) + } + }) + } +} + +// TestSSE405AllowHeader verifies RFC 9110 §15.5.6 compliance: +// 405 Method Not Allowed responses MUST include an Allow header. +func TestSSE405AllowHeader(t *testing.T) { + server := NewServer(testImpl, nil) + + handler := NewSSEHandler(func(req *http.Request) *Server { return server }, nil) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + methods := []string{"PUT", "PATCH", "DELETE", "OPTIONS"} + for _, method := range methods { + t.Run(method, func(t *testing.T) { + req, err := http.NewRequest(method, httpServer.URL, nil) + if err != nil { + t.Fatal(err) + } + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + if got, want := resp.StatusCode, http.StatusMethodNotAllowed; got != want { + t.Errorf("status code: got %d, want %d", got, want) + } + + allow := resp.Header.Get("Allow") + if allow != "GET, POST" { + t.Errorf("Allow header: got %q, want %q", allow, "GET, POST") + } + }) + } +} + +// TestSSELocalhostProtection verifies that DNS rebinding protection +// is automatically enabled for localhost servers. +func TestSSELocalhostProtection(t *testing.T) { + server := NewServer(testImpl, nil) + + tests := []struct { + name string + listenAddr string + hostHeader string + disableProtection bool + wantStatus int + }{ + { + name: "127.0.0.1 accepts 127.0.0.1", + listenAddr: "127.0.0.1:0", + hostHeader: "127.0.0.1:1234", + wantStatus: http.StatusOK, + }, + { + name: "127.0.0.1 accepts localhost", + listenAddr: "127.0.0.1:0", + hostHeader: "localhost:1234", + wantStatus: http.StatusOK, + }, + { + name: "127.0.0.1 rejects evil.com", + listenAddr: "127.0.0.1:0", + hostHeader: "evil.com", + wantStatus: http.StatusForbidden, + }, + { + name: "127.0.0.1 rejects evil.com:80", + listenAddr: "127.0.0.1:0", + hostHeader: "evil.com:80", + wantStatus: http.StatusForbidden, + }, + { + name: "127.0.0.1 rejects localhost.evil.com", + listenAddr: "127.0.0.1:0", + hostHeader: "localhost.evil.com", + wantStatus: http.StatusForbidden, + }, + { + name: "0.0.0.0 via localhost rejects evil.com", + listenAddr: "0.0.0.0:0", + hostHeader: "evil.com", + wantStatus: http.StatusForbidden, + }, + { + name: "disabled accepts evil.com", + listenAddr: "127.0.0.1:0", + hostHeader: "evil.com", + disableProtection: true, + wantStatus: http.StatusOK, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + opts := &SSEOptions{ + DisableLocalhostProtection: tt.disableProtection, + } + handler := NewSSEHandler(func(req *http.Request) *Server { return server }, opts) + + listener, err := net.Listen("tcp", tt.listenAddr) + if err != nil { + t.Fatalf("Failed to listen on %s: %v", tt.listenAddr, err) + } + defer listener.Close() + + srv := &http.Server{Handler: handler} + go srv.Serve(listener) + defer srv.Close() + + // Use a GET request since it's the entry point for SSE sessions. + // For accepted requests, the response will be a hanging SSE stream, + // but we only need to check the initial status code. + req, err := http.NewRequest("GET", fmt.Sprintf("http://%s", listener.Addr().String()), nil) + if err != nil { + t.Fatal(err) + } + req.Host = tt.hostHeader + req.Header.Set("Accept", "text/event-stream") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + if got := resp.StatusCode; got != tt.wantStatus { + t.Errorf("Status code: got %d, want %d", got, tt.wantStatus) + } + }) + } +} + +// TestSSEServerTransportMaxRequestBodyBytes exercises the limit on the exported +// [SSEServerTransport] directly, since it is documented as usable on its own. +func TestSSEServerTransportMaxRequestBodyBytes(t *testing.T) { + requestOfSize := func(n int) []byte { + b := bytes.Repeat([]byte(" "), n) + b[0], b[n-1] = '{', '}' + return b + } + testCases := []struct { + name string + limit int64 + request []byte + wantCode int + }{ + { + name: "over limit returns", + limit: 1024, + request: requestOfSize(2048), + wantCode: http.StatusRequestEntityTooLarge, + }, + { + name: "under limit accepted", + limit: 1024, + request: []byte(`{"jsonrpc":"2.0","id":1,"method":"ping"}`), + wantCode: http.StatusAccepted, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + tr := &SSEServerTransport{ + Endpoint: "/msg", + Response: httptest.NewRecorder(), + MaxRequestBodyBytes: tc.limit, + } + if _, err := tr.Connect(context.Background()); err != nil { + t.Fatalf("Connect: %v", err) + } + req := httptest.NewRequest(http.MethodPost, "/msg", bytes.NewReader(tc.request)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + tr.ServeHTTP(rec, req) + if rec.Code != tc.wantCode { + t.Errorf("status = %d, want %d", rec.Code, tc.wantCode) + } + }) + } +} + +func TestSSEClientTransportBufferedEvents(t *testing.T) { + const endpoint = "event: endpoint\ndata: /messages\n\n" + const message = "event: message\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/tools/list_changed\"}\n\n" + for _, coalesced := range []bool{false, true} { + t.Run(fmt.Sprintf("coalesced=%t", coalesced), func(t *testing.T) { + reader, writer := io.Pipe() + defer reader.Close() + defer writer.Close() + prefix := endpoint + if coalesced { + prefix += message + } + transport := &SSEClientTransport{ + Endpoint: "http://example.com/sse", + HTTPClient: &http.Client{Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + Body: struct { + io.Reader + io.Closer + }{io.MultiReader(strings.NewReader(prefix), reader), reader}, + }, nil + })}, + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + conn, err := transport.Connect(ctx) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + if !coalesced { + go func() { _, _ = io.WriteString(writer, message) }() + } + msg, err := conn.Read(ctx) + if err != nil { + t.Fatalf("Read: %v", err) + } + if req, ok := msg.(*jsonrpc.Request); !ok || req.Method != "notifications/tools/list_changed" { + t.Fatalf("Read = %v, want tools/list_changed notification", msg) + } + }) + } +} diff --git a/mcp/streamable.go b/mcp/streamable.go index d8ce45e..f24badb 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -2,11 +2,16 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. +// NOTE: see streamable_server.go and streamable_client.go for detailed +// documentation of the streamable server design. +// TODO: move the client and server logic into those files. + package mcp import ( "bytes" "context" + crand "crypto/rand" "encoding/json" "errors" "fmt" @@ -15,6 +20,8 @@ import ( "maps" "math" "math/rand/v2" + "mime" + "net" "net/http" "slices" "strconv" @@ -24,13 +31,13 @@ import ( "time" "github.com/modelcontextprotocol/go-sdk/auth" + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/internal/mcpgodebug" + "github.com/modelcontextprotocol/go-sdk/internal/util" + "github.com/modelcontextprotocol/go-sdk/internal/xcontext" "github.com/modelcontextprotocol/go-sdk/jsonrpc" -) - -const ( - protocolVersionHeader = "Mcp-Protocol-Version" - sessionIDHeader = "Mcp-Session-Id" + "golang.org/x/oauth2" ) // A StreamableHTTPHandler is an http.Handler that serves streamable MCP @@ -50,6 +57,10 @@ type StreamableHTTPHandler struct { type sessionInfo struct { session *ServerSession transport *StreamableServerTransport + // userID is the user ID from the TokenInfo when the session was created. + // If non-empty, subsequent requests must have the same user ID to prevent + // session hijacking. + userID string // If timeout is set, automatically close the session after an idle period. timeout time.Duration @@ -82,7 +93,7 @@ func (i *sessionInfo) startPOST() { i.refs++ } -// endPOST sigals that a request for this session is ending, starting the +// endPOST signals that a request for this session is ending, starting the // timeout if there are no other requests running. func (i *sessionInfo) endPOST() { if i.timeout <= 0 { @@ -117,12 +128,21 @@ func (i *sessionInfo) stopTimer() { type StreamableHTTPOptions struct { // Stateless controls whether the session is 'stateless'. // - // A stateless server does not validate the Mcp-Session-Id header, and uses a - // temporary session with default initialization parameters. Any + // A stateless server does not read or set the Mcp-Session-Id header, and + // uses a temporary session with default initialization parameters for each + // request. [ServerOptions.GetSessionID] is not consulted. Any // server->client request is rejected immediately as there's no way for the // client to respond. Server->Client notifications may reach the client if // they are made in the context of an incoming request, as described in the // documentation for [StreamableServerTransport]. + // In Stateless mode, GET and DELETE requests return 405 Method Not Allowed. + // + // This mode aligns with the sessionless direction of the MCP spec; see + // [SEP-2567]. The previous behavior, in which stateless servers still + // honored Mcp-Session-Id, can be restored temporarily via the + // MCPGODEBUG compatibility parameter "allowsessionsinstateless=1". + // + // [SEP-2567]: https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2567 Stateless bool // TODO(#148): support session retention (?) @@ -150,8 +170,59 @@ type StreamableHTTPOptions struct { // // If SessionTimeout is the zero value, idle sessions are never closed. SessionTimeout time.Duration + + // DisableLocalhostProtection disables automatic DNS rebinding protection. + // By default, requests arriving via a localhost address (127.0.0.1, [::1]) + // that have a non-localhost Host header are rejected with 403 Forbidden. + // This protects against DNS rebinding attacks regardless of whether the + // server is listening on localhost specifically or on 0.0.0.0. + // + // Only disable this if you understand the security implications. + // See: https://modelcontextprotocol.io/specification/2025-11-25/basic/security_best_practices#local-mcp-server-compromise + DisableLocalhostProtection bool + + // CrossOriginProtection allows to customize cross-origin protection. + // The deny handler set in the CrossOriginProtection through SetDenyHandler + // is ignored. + // If nil, no cross-origin protection is applied. + // + // Deprecated: wrap the handler with cross-origin protection middleware + // instead. For example: + // + // handler := mcp.NewStreamableHTTPHandler(...) + // protection := http.NewCrossOriginProtection() + // protectedHandler := protection.Handler(handler) + CrossOriginProtection *http.CrossOriginProtection + + // MaxRequestBodyBytes limits the number of bytes read from any incoming + // HTTP request body. Requests that exceed this limit are rejected with + // 413 Request Entity Too Large. + // + // The limit is enforced during the read, so it applies uniformly to + // requests using Content-Length, Transfer-Encoding: chunked, or HTTP/2 + // (which has no Content-Length). + // + // If zero, [DefaultMaxRequestBodyBytes] is used. + // A negative value disables the limit entirely; do not use this on + // servers exposed to untrusted clients. + MaxRequestBodyBytes int64 + + // PropagateRequestCancellation, when true, ties the in-flight handler + // context to the originating HTTP request's context. Only applies to + // requests using the >= 2026-07-28 protocol, where the POST is the whole + // request lifecycle. + // The handler context cancels whenever the HTTP request context does as the + // response can no longer be delivered, so cancelling the handler is safe. + // + // Requests using older protocol versions (including those routed through + // the allowsessionsinstateless compatibility path) are unaffected. + PropagateRequestCancellation bool } +// DefaultMaxRequestBodyBytes is the default value used for +// [StreamableHTTPOptions.MaxRequestBodyBytes] when it is left at zero. +const DefaultMaxRequestBodyBytes = 4 << 20 // 4 MiB + // NewStreamableHTTPHandler returns a new [StreamableHTTPHandler]. // // The getServer function is used to create or look up servers for new @@ -166,8 +237,10 @@ func NewStreamableHTTPHandler(getServer func(*http.Request) *Server, opts *Strea h.opts = *opts } - if h.opts.Logger == nil { // ensure we have a logger - h.opts.Logger = ensureLogger(nil) + h.opts.Logger = ensureLogger(h.opts.Logger) + + if h.opts.MaxRequestBodyBytes == 0 { + h.opts.MaxRequestBodyBytes = DefaultMaxRequestBodyBytes } return h @@ -196,243 +269,502 @@ func (h *StreamableHTTPHandler) closeAll() { } } +// allowsessionsinstateless is a compatibility parameter that restores the old +// behavior of reading and using Mcp-Session-Id headers in stateless mode. When +// set to "1", stateless servers will read the session ID from the request +// header (or generate one via GetSessionID), set it on response headers, and +// accept DELETE requests. When unset (the default), stateless servers ignore +// session IDs entirely and reject DELETE with 405. +// See the documentation for the mcpgodebug package for instructions how to enable it. +// The option will be removed in the 1.9.0 version of the SDK. +var allowsessionsinstateless = mcpgodebug.Value("allowsessionsinstateless") + +// noprotocolerrorbody is a compatibility parameter that restores the previous +// behavior of [streamableClientConn.checkResponse]. When unset (the default), +// the client attempts to decode the response body of a non-2xx response as a +// JSON-RPC error. If successful, the underlying JSON-RPC error is surfaced to +// the caller and wrapped with [jsonrpc2.ErrRejected] so that per-call +// rejections do not tear down the session. When set to "1", the client +// ignores response bodies on non-2xx responses and any non-transient error +// permanently fails the connection. +var noprotocolerrorbody = mcpgodebug.Value("noprotocolerrorbody") + +// plaintextstatefulrejection is a compatibility parameter that restores the +// previous behavior of a stateful [StreamableHTTPHandler] when it receives a +// request carrying SEP-2575 per-request metadata (i.e. an +// `io.modelcontextprotocol/protocolVersion` `_meta` field, or an +// `MCP-Protocol-Version` header >= 2026-07-28). When unset (the default), the +// server responds with a JSON-RPC error of code +// [CodeUnsupportedProtocolVersion] and an [UnsupportedProtocolVersionData] +// payload listing its supported legacy versions, as required by SEP-2575. When +// set to "1", the server instead responds with a plain-text `http.Error` body +// mentioning the `StreamableHTTPOptions.Stateless` field. +var plaintextstatefulrejection = mcpgodebug.Value("plaintextstatefulrejection") + +// writeJSONRPCError writes a JSON-RPC error response with the given HTTP +// status code, request ID (may be a zero ID for errors that occur before the +// request body has been parsed), and JSON-RPC error. +func writeJSONRPCError(w http.ResponseWriter, status int, id jsonrpc.ID, jerr *jsonrpc.Error) { + resp := &jsonrpc.Response{ID: id, Error: jerr} + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + if data, err := jsonrpc2.EncodeMessage(resp); err == nil { + w.Write(data) + } +} + func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { - // Allow multiple 'Accept' headers. - // https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Accept#syntax - accept := strings.Split(strings.Join(req.Header.Values("Accept"), ","), ",") - var jsonOK, streamOK bool - for _, c := range accept { - switch strings.TrimSpace(c) { - case "application/json", "application/*": - jsonOK = true - case "text/event-stream", "text/*": - streamOK = true - case "*/*": - jsonOK = true - streamOK = true + // DNS rebinding protection: auto-enabled for localhost servers. + // See: https://modelcontextprotocol.io/specification/2025-11-25/basic/security_best_practices#local-mcp-server-compromise + if !h.opts.DisableLocalhostProtection { + if localAddr, ok := req.Context().Value(http.LocalAddrContextKey).(net.Addr); ok && localAddr != nil { + if util.IsLoopback(localAddr.String()) && !util.IsLoopback(req.Host) { + http.Error(w, fmt.Sprintf("Forbidden: invalid Host header %q", req.Host), http.StatusForbidden) + return + } } } - if req.Method == http.MethodGet { - if !streamOK { - http.Error(w, "Accept must contain 'text/event-stream' for GET requests", http.StatusBadRequest) + if h.opts.CrossOriginProtection != nil { + if err := h.opts.CrossOriginProtection.Check(req); err != nil { + http.Error(w, err.Error(), http.StatusForbidden) return } - } else if (!jsonOK || !streamOK) && req.Method != http.MethodDelete { // TODO: consolidate with handling of http method below. + } + + // Bound the request body to protect against OOM attacks. + if req.Body != nil && h.opts.MaxRequestBodyBytes > 0 { + req.Body = http.MaxBytesReader(w, req.Body, h.opts.MaxRequestBodyBytes) + } + + // [§2.7] of the spec (2025-06-18): validate the MCP-Protocol-Version + // header. If provided, it must be a version the server negotiates, which + // [ServerOptions.SupportedProtocolVersions] may have narrowed. If absent, + // the version is unknown (the request may be an initialize for any + // version). + // + // [§2.7]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#protocol-version-header + protocolVersion := req.Header.Get(protocolVersionHeader) + var supportedVersions []string + if server := h.getServer(req); server != nil { + supportedVersions = server.protocolVersions + } else { + supportedVersions = supportedProtocolVersions + } + if protocolVersion != "" && !slices.Contains(supportedVersions, protocolVersion) && protocolVersion < protocolVersion20260728 { + http.Error(w, fmt.Sprintf("Bad Request: Unsupported protocol version (supported versions: %s)", strings.Join(supportedVersions, ",")), http.StatusBadRequest) + return + } + req = req.WithContext(context.WithValue(req.Context(), protocolVersionContextKey{}, protocolVersion)) + + if h.opts.Stateless { + h.serveStateless(w, req) + } else { + h.serveStateful(w, req) + } +} + +// serveStateless handles requests for stateless servers. +// Stateless servers only support POST. Each request creates a temporary +// session that is closed when the request completes. +// +// When the allowsessionsinstateless compatibility flag is set, DELETE is also +// accepted (as a no-op) and session IDs are read from the request header. +func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http.Request) { + legacySessions := allowsessionsinstateless == "1" + + if req.Method == http.MethodDelete && legacySessions { + h.serveStatelessLegacyDELETE(w, req) + return + } + + if req.Method != http.MethodPost { + // RFC 9110 §15.5.6: 405 responses MUST include Allow header. + w.Header().Set("Allow", "POST") + http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed) + return + } + + if baseMediaType(req.Header.Get("Content-Type")) != "application/json" { + http.Error(w, "Content-Type must be 'application/json'", http.StatusUnsupportedMediaType) + return + } + + // Accept must contain both 'application/json' and 'text/event-stream'. + jsonOK, streamOK := streamableAccepts(req.Header.Values("Accept")) + if !jsonOK || !streamOK { + http.Error(w, "Accept must contain both 'application/json' and 'text/event-stream'", http.StatusBadRequest) + return + } + + server := h.getServer(req) + if server == nil { + http.Error(w, "no server available", http.StatusBadRequest) + return + } + + info, err := h.ephemeralConnectOpts(req) + if err != nil { + var mbe *http.MaxBytesError + if errors.As(err, &mbe) { + http.Error(w, fmt.Sprintf("request body exceeds %d bytes", mbe.Limit), http.StatusRequestEntityTooLarge) + return + } + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + var sessionID string + if legacySessions && !info.usesNewProtocol { + sessionID = req.Header.Get(sessionIDHeader) + if sessionID == "" { + sessionID = server.opts.GetSessionID() + } + } + + transport := &StreamableServerTransport{ + SessionID: sessionID, + Stateless: true, + EventStore: h.opts.EventStore, + jsonResponse: h.opts.JSONResponse, + logger: h.opts.Logger, + shouldPropagateCancellation: info.usesNewProtocol && (info.isSubscriptionsListen || h.opts.PropagateRequestCancellation), + } + + session, err := connectStreamable(req.Context(), server, transport, info.opts) + if err != nil { + h.opts.Logger.Error(fmt.Sprintf("failed to connect: %v", err)) + http.Error(w, "failed connection", http.StatusInternalServerError) + return + } + defer session.Close() + + transport.ServeHTTP(w, req) +} + +// serveStatelessLegacyDELETE handles DELETE requests in stateless mode when the +// allowsessionsinstateless compatibility flag is set. DELETE requires a +// Mcp-Session-Id header but is otherwise a no-op since stateless servers don't +// persist sessions. +func (h *StreamableHTTPHandler) serveStatelessLegacyDELETE(w http.ResponseWriter, req *http.Request) { + sessionID := req.Header.Get(sessionIDHeader) + if sessionID == "" { + http.Error(w, "Bad Request: DELETE requires an Mcp-Session-Id header", http.StatusBadRequest) + return + } + w.WriteHeader(http.StatusNoContent) +} + +type ephemeralConnectInfo struct { + opts *ServerSessionOptions + usesNewProtocol bool + isSubscriptionsListen bool +} + +// ephemeralConnectOpts peeks at the request body to determine connection +// parameters and whether protocol version >= 2026-06-30 (SEP-2575). +// +// For old-protocol requests, default session state is synthesized so that +// the session's init gate doesn't reject the request. +// +// It is used for both stateless servers and stateful servers with no session ID. +func (h *StreamableHTTPHandler) ephemeralConnectOpts(req *http.Request) (*ephemeralConnectInfo, error) { + protocolVersion := protocolVersionFromContext(req.Context()) + if protocolVersion == "" { + protocolVersion = protocolVersion20250326 + } + + var hasInitialize, hasInitialized, usesNewProtocol, isSubscriptionsListen bool + body, err := io.ReadAll(req.Body) + if err != nil { + // Preserve *http.MaxBytesError so serveStateless can respond with 413. + return nil, fmt.Errorf("failed to read body: %w", err) + } + req.Body.Close() + req.Body = io.NopCloser(bytes.NewBuffer(body)) + msgs, _, err := readBatch(body) + if err == nil { + for _, msg := range msgs { + if r, ok := msg.(*jsonrpc.Request); ok { + switch r.Method { + case methodInitialize: + hasInitialize = true + case notificationInitialized: + hasInitialized = true + case methodSubscriptionsListen: + isSubscriptionsListen = true + } + if protocolVersion >= protocolVersion20260728 { + usesNewProtocol = true + } + } + } + } + + state := new(ServerSessionState) + // Only synthesize fake InitializeParams/InitializedParams for old-protocol + // requests. + if !hasInitialize && !usesNewProtocol { + state.InitializeParams = &InitializeParams{ + ProtocolVersion: protocolVersion, + } + } + if !hasInitialized && !usesNewProtocol { + state.InitializedParams = new(InitializedParams) + } + if !usesNewProtocol { + state.LogLevel = "info" + } + return &ephemeralConnectInfo{ + opts: &ServerSessionOptions{ + State: state, + }, + usesNewProtocol: usesNewProtocol, + isSubscriptionsListen: isSubscriptionsListen, + }, nil +} + +func connectStreamable(ctx context.Context, server *Server, transport *StreamableServerTransport, opts *ServerSessionOptions) (*ServerSession, error) { + s, err := server.Connect(ctx, transport, opts) + if err != nil { + return nil, err + } + transport.connection.server = server + transport.connection.toolLookup = server.getServerTool + return s, nil +} + +// serveStateful handles requests for stateful servers. +// Stateful servers support GET, POST, and DELETE, and maintain persistent +// sessions keyed by session ID. +func (h *StreamableHTTPHandler) serveStateful(w http.ResponseWriter, req *http.Request) { + switch req.Method { + case http.MethodGet: + h.serveStatefulGET(w, req) + case http.MethodPost: + h.serveStatefulPOST(w, req) + case http.MethodDelete: + h.serveStatefulDELETE(w, req) + default: + // RFC 9110 §15.5.6: 405 responses MUST include Allow header. + w.Header().Set("Allow", "GET, POST, DELETE") + http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed) + } +} + +// lookupSession looks up a session by the Mcp-Session-Id header value. +// It returns the session info and whether the caller should proceed. If ok is +// false, an error response has been written. The sessionID must be non-empty; +// callers are responsible for checking this before calling lookupSession. +func (h *StreamableHTTPHandler) lookupSession(w http.ResponseWriter, req *http.Request, sessionID string) (info *sessionInfo, ok bool) { + h.mu.Lock() + info = h.sessions[sessionID] + h.mu.Unlock() + if info == nil { + http.Error(w, "session not found", http.StatusNotFound) + return nil, false + } + if info.userID != "" { + tokenInfo := auth.TokenInfoFromContext(req.Context()) + if tokenInfo == nil || tokenInfo.UserID != info.userID { + http.Error(w, "session user mismatch", http.StatusForbidden) + return nil, false + } + } + return info, true +} + +// serveStatefulGET handles GET requests for standalone SSE streams. +// GET requires a valid Mcp-Session-Id header. +func (h *StreamableHTTPHandler) serveStatefulGET(w http.ResponseWriter, req *http.Request) { + if _, streamOK := streamableAccepts(req.Header.Values("Accept")); !streamOK { + http.Error(w, "Accept must contain 'text/event-stream' for GET requests", http.StatusBadRequest) + return + } + + sessionID := req.Header.Get(sessionIDHeader) + if sessionID == "" { + http.Error(w, "Bad Request: GET requires an Mcp-Session-Id header", http.StatusBadRequest) + return + } + + sessInfo, ok := h.lookupSession(w, req, sessionID) + if !ok { + return + } + + sessInfo.transport.ServeHTTP(w, req) +} + +// serveStatefulDELETE handles DELETE requests for session termination. +// DELETE requires a valid Mcp-Session-Id header. +func (h *StreamableHTTPHandler) serveStatefulDELETE(w http.ResponseWriter, req *http.Request) { + sessionID := req.Header.Get(sessionIDHeader) + if sessionID == "" { + http.Error(w, "Bad Request: DELETE requires an Mcp-Session-Id header", http.StatusBadRequest) + return + } + + sessInfo, ok := h.lookupSession(w, req, sessionID) + if !ok { + return + } + + sessInfo.session.Close() + w.WriteHeader(http.StatusNoContent) +} + +// serveStatefulPOST handles POST requests for stateful servers. +// POST may arrive with or without a Mcp-Session-Id header. Without a session +// ID, a new session is created (this is the normal path for the first +// initialize request). +func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *http.Request) { + if baseMediaType(req.Header.Get("Content-Type")) != "application/json" { + http.Error(w, "Content-Type must be 'application/json'", http.StatusUnsupportedMediaType) + return + } + + jsonOK, streamOK := streamableAccepts(req.Header.Values("Accept")) + if !jsonOK || !streamOK { http.Error(w, "Accept must contain both 'application/json' and 'text/event-stream'", http.StatusBadRequest) return } sessionID := req.Header.Get(sessionIDHeader) - var sessInfo *sessionInfo + + // Look up existing session if a session ID was provided. if sessionID != "" { - h.mu.Lock() - sessInfo = h.sessions[sessionID] - h.mu.Unlock() - if sessInfo == nil && !h.opts.Stateless { - // Unless we're in 'stateless' mode, which doesn't perform any Session-ID - // validation, we require that the session ID matches a known session. - // - // In stateless mode, a temporary transport is be created below. - http.Error(w, "session not found", http.StatusNotFound) + sessInfo, ok := h.lookupSession(w, req, sessionID) + if !ok { return } - } - - if req.Method == http.MethodDelete { - if sessionID == "" { - http.Error(w, "Bad Request: DELETE requires an Mcp-Session-Id header", http.StatusBadRequest) - return - } - if sessInfo != nil { // sessInfo may be nil in stateless mode - // Closing the session also removes it from h.sessions, due to the - // onClose callback. - sessInfo.session.Close() - } - w.WriteHeader(http.StatusNoContent) + sessInfo.startPOST() + defer sessInfo.endPOST() + sessInfo.transport.ServeHTTP(w, req) return } - switch req.Method { - case http.MethodPost, http.MethodGet: - if req.Method == http.MethodGet && (h.opts.Stateless || sessionID == "") { - http.Error(w, "GET requires an active session", http.StatusMethodNotAllowed) - return - } - default: - w.Header().Set("Allow", "GET, POST, DELETE") - http.Error(w, "Method Not Allowed: streamable MCP servers support GET, POST, and DELETE requests", http.StatusMethodNotAllowed) + // No session ID: create a new session. + server := h.getServer(req) + if server == nil { + http.Error(w, "no server available", http.StatusBadRequest) return } + sessionID = server.opts.GetSessionID() - // [§2.7] of the spec (2025-06-18) states: - // - // "If using HTTP, the client MUST include the MCP-Protocol-Version: - // HTTP header on all subsequent requests to the MCP - // server, allowing the MCP server to respond based on the MCP protocol - // version. - // - // For example: MCP-Protocol-Version: 2025-06-18 - // The protocol version sent by the client SHOULD be the one negotiated during - // initialization. - // - // For backwards compatibility, if the server does not receive an - // MCP-Protocol-Version header, and has no other way to identify the version - - // for example, by relying on the protocol version negotiated during - // initialization - the server SHOULD assume protocol version 2025-03-26. - // - // If the server receives a request with an invalid or unsupported - // MCP-Protocol-Version, it MUST respond with 400 Bad Request." - // - // Since this wasn't present in the 2025-03-26 version of the spec, this - // effectively means: - // 1. IF the client provides a version header, it must be a supported - // version. - // 2. In stateless mode, where we've lost the state of the initialize - // request, we assume that whatever the client tells us is the truth (or - // assume 2025-03-26 if the client doesn't say anything). - // - // This logic matches the typescript SDK. - // - // [§2.7]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#protocol-version-header - protocolVersion := req.Header.Get(protocolVersionHeader) - if protocolVersion == "" { - protocolVersion = protocolVersion20250326 - } - if !slices.Contains(supportedProtocolVersions, protocolVersion) { - http.Error(w, fmt.Sprintf("Bad Request: Unsupported protocol version (supported versions: %s)", strings.Join(supportedProtocolVersions, ",")), http.StatusBadRequest) - return + transport := &StreamableServerTransport{ + SessionID: sessionID, + Stateless: false, + EventStore: h.opts.EventStore, + jsonResponse: h.opts.JSONResponse, + logger: h.opts.Logger, } - if sessInfo == nil { - server := h.getServer(req) - if server == nil { - // The getServer argument to NewStreamableHTTPHandler returned nil. - http.Error(w, "no server available", http.StatusBadRequest) - return - } - if sessionID == "" { - // In stateless mode, sessionID may be nonempty even if there's no - // existing transport. - sessionID = server.opts.GetSessionID() - } - transport := &StreamableServerTransport{ - SessionID: sessionID, - Stateless: h.opts.Stateless, - EventStore: h.opts.EventStore, - jsonResponse: h.opts.JSONResponse, - logger: h.opts.Logger, - } - - // To support stateless mode, we initialize the session with a default - // state, so that it doesn't reject subsequent requests. - var connectOpts *ServerSessionOptions - if h.opts.Stateless { - // Peek at the body to see if it is initialize or initialized. - // We want those to be handled as usual. - var hasInitialize, hasInitialized bool - { - // TODO: verify that this allows protocol version negotiation for - // stateless servers. - body, err := io.ReadAll(req.Body) - if err != nil { - http.Error(w, "failed to read body", http.StatusInternalServerError) - return - } - req.Body.Close() - - // Reset the body so that it can be read later. - req.Body = io.NopCloser(bytes.NewBuffer(body)) - - msgs, _, err := readBatch(body) - if err == nil { - for _, msg := range msgs { - if req, ok := msg.(*jsonrpc.Request); ok { - switch req.Method { - case methodInitialize: - hasInitialize = true - case notificationInitialized: - hasInitialized = true - } - } - } - } - } - - // If we don't have InitializeParams or InitializedParams in the request, - // set the initial state to a default value. - state := new(ServerSessionState) - if !hasInitialize { - state.InitializeParams = &InitializeParams{ - ProtocolVersion: protocolVersion, - } - } - if !hasInitialized { - state.InitializedParams = new(InitializedParams) - } - state.LogLevel = "info" - connectOpts = &ServerSessionOptions{ - State: state, - } - } else { - // Cleanup is only required in stateful mode, as transportation is - // not stored in the map otherwise. - connectOpts = &ServerSessionOptions{ - onClose: func() { - h.mu.Lock() - defer h.mu.Unlock() - if info, ok := h.sessions[transport.SessionID]; ok { - info.stopTimer() - delete(h.sessions, transport.SessionID) - if h.onTransportDeletion != nil { - h.onTransportDeletion(transport.SessionID) - } - } - }, - } - } - - // Pass req.Context() here, to allow middleware to add context values. - // The context is detached in the jsonrpc2 library when handling the - // long-running stream. - session, err := server.Connect(req.Context(), transport, connectOpts) + // Sessions without a session ID (GetSessionID returned "") are ephemeral: + // there's no way to address them, so they are closed after the request. + // This can happen when ServerOptions.GetSessionID is explicitly set to + // return "" to suppress session IDs entirely. It also covers any request + // that arrives before a session exists (e.g. initialize or ping) on a + // server configured this way. + if sessionID == "" { + info, err := h.ephemeralConnectOpts(req) if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + session, err := connectStreamable(req.Context(), server, transport, info.opts) + if err != nil { + h.opts.Logger.Error(fmt.Sprintf("failed to connect: %v", err)) http.Error(w, "failed connection", http.StatusInternalServerError) return } - sessInfo = &sessionInfo{ - session: session, - transport: transport, - } + defer session.Close() + transport.ServeHTTP(w, req) + return + } - if h.opts.Stateless { - // Stateless mode: close the session when the request exits. - defer session.Close() // close the fake session after handling the request - } else { - // Otherwise, save the transport so that it can be reused - - // Clean up the session when it times out. - // - // Note that the timer here may fire multiple times, but - // sessInfo.session.Close is idempotent. - if h.opts.SessionTimeout > 0 { - sessInfo.timeout = h.opts.SessionTimeout - sessInfo.timer = time.AfterFunc(sessInfo.timeout, func() { - sessInfo.session.Close() - }) - } + connectOpts := &ServerSessionOptions{ + onClose: func() { h.mu.Lock() - h.sessions[transport.SessionID] = sessInfo - h.mu.Unlock() + defer h.mu.Unlock() + if info, ok := h.sessions[transport.SessionID]; ok { + info.stopTimer() + delete(h.sessions, transport.SessionID) + if h.onTransportDeletion != nil { + h.onTransportDeletion(transport.SessionID) + } + } + }, + } + + // Pass req.Context() here, to allow middleware to add context values. + // The context is detached in the jsonrpc2 library when handling the + // long-running stream. + session, err := connectStreamable(req.Context(), server, transport, connectOpts) + if err != nil { + h.opts.Logger.Error(fmt.Sprintf("failed to connect: %v", err)) + http.Error(w, "failed connection", http.StatusInternalServerError) + return + } + // Capture the user ID from the token info to enable session hijacking + // prevention on subsequent requests. + var userID string + if tokenInfo := auth.TokenInfoFromContext(req.Context()); tokenInfo != nil { + userID = tokenInfo.UserID + } + sessInfo := &sessionInfo{ + session: session, + transport: transport, + userID: userID, + } + + if h.opts.SessionTimeout > 0 { + sessInfo.timeout = h.opts.SessionTimeout + sessInfo.timer = time.AfterFunc(sessInfo.timeout, func() { + sessInfo.session.Close() + }) + } + h.mu.Lock() + h.sessions[transport.SessionID] = sessInfo + h.mu.Unlock() + defer func() { + // If initialization failed, clean up the session (#578). + if session.InitializeParams() == nil { + session.Close() + } + }() + + sessInfo.startPOST() + defer sessInfo.endPOST() + sessInfo.transport.ServeHTTP(w, req) +} + +func streamableAccepts(values []string) (jsonOK, streamOK bool) { + for _, value := range values { + for _, raw := range strings.Split(value, ",") { + token := strings.TrimSpace(raw) + // Ignore Accept parameters like ";charset=utf-8"; match the base media type. + base, _, _ := strings.Cut(token, ";") + switch strings.ToLower(strings.TrimSpace(base)) { + case "application/json", "application/*": + jsonOK = true + case "text/event-stream", "text/*": + streamOK = true + case "*/*": + jsonOK = true + streamOK = true + } } } + return jsonOK, streamOK +} - if req.Method == http.MethodPost { - sessInfo.startPOST() - defer sessInfo.endPOST() +func baseMediaType(value string) string { + mediaType, _, err := mime.ParseMediaType(value) + if err != nil { + return "" } - - sessInfo.transport.ServeHTTP(w, req) + return mediaType } // A StreamableServerTransport implements the server side of the MCP streamable @@ -496,6 +828,10 @@ type StreamableServerTransport struct { // to write their own streamable HTTP handler. logger *slog.Logger + // shouldPropagateCancellation is forwarded to the underlying + // [streamableServerConn]. See its docstring. + shouldPropagateCancellation bool + // connection is non-nil if and only if the transport has been connected. connection *streamableServerConn } @@ -506,15 +842,16 @@ func (t *StreamableServerTransport) Connect(ctx context.Context) (Connection, er return nil, fmt.Errorf("transport already connected") } t.connection = &streamableServerConn{ - sessionID: t.SessionID, - stateless: t.Stateless, - eventStore: t.EventStore, - jsonResponse: t.jsonResponse, - logger: ensureLogger(t.logger), // see #556: must be non-nil - incoming: make(chan jsonrpc.Message, 10), - done: make(chan struct{}), - streams: make(map[string]*stream), - requestStreams: make(map[jsonrpc.ID]string), + sessionID: t.SessionID, + stateless: t.Stateless, + eventStore: t.EventStore, + jsonResponse: t.jsonResponse, + logger: ensureLogger(t.logger), // see #556: must be non-nil + shouldPropagateCancellation: t.shouldPropagateCancellation, + incoming: make(chan jsonrpc.Message, 10), + done: make(chan struct{}), + streams: make(map[string]*stream), + requestStreams: make(map[jsonrpc.ID]string), } // Stream 0 corresponds to the standalone SSE stream. // @@ -528,14 +865,34 @@ func (t *StreamableServerTransport) Connect(ctx context.Context) (Connection, er return t.connection, nil } +// The streamable HTTP transport supports every legacy SDK protocol version, +// but the SEP-2575 >= 2026-07-28 protocol is only supported when the +// transport is configured as stateless. +func (t *StreamableServerTransport) SupportsProtocolVersion(version string) bool { + if version >= protocolVersion20260728 { + return t.Stateless && slices.Contains(supportedProtocolVersions, version) + } + return slices.Contains(supportedProtocolVersions, version) +} + type streamableServerConn struct { sessionID string stateless bool jsonResponse bool eventStore EventStore + // shouldPropagateCancellation is true when the underlying HTTP request's + // lifetime IS the connection's cancellation signal (e.g., a stateless + // POST that owns a long-lived subscriptions/listen stream). It is read + // by the [cancellationPropagator] interface so the jsonrpc2 layer wires + // handler contexts to observe the carrier's cancellation. + shouldPropagateCancellation bool + logger *slog.Logger + server *Server + toolLookup func(name string) (*serverTool, bool) + incoming chan jsonrpc.Message // messages from the client to the server mu sync.Mutex // guards all fields below @@ -569,6 +926,15 @@ func (c *streamableServerConn) SessionID() string { return c.sessionID } +// propagateCancellation implements [cancellationPropagator]. It returns true +// when this connection is bound to a single HTTP request whose lifetime +// should drive request-handler cancellation — for example, a stateless POST +// carrying a long-lived subscriptions/listen stream that must unwind when +// the client TCP-disconnects. +func (c *streamableServerConn) propagateCancellation() bool { + return c.shouldPropagateCancellation +} + // A stream is a single logical stream of SSE events within a server session. // A stream begins with a client request, or with a client GET that has // no Last-Event-ID header. @@ -582,31 +948,210 @@ type stream struct { // The standalone SSE stream has id "". id string + // logger is used for logging errors during stream operations. + logger *slog.Logger + // mu guards the fields below, as well as storage of new messages in the // connection's event store (if any). mu sync.Mutex - // If non-nil, deliver writes data directly to the HTTP response. + // If pendingJSONMessages is non-nil, this is a JSON stream and messages are + // collected here until the stream is complete, at which point they are + // flushed as a single JSON response. Note that the non-nilness of this field + // is significant, as it signals the expected content type. // - // Only one HTTP response may receive messages at a given time. An active - // HTTP connection acquires ownership of the stream by setting this field. - deliver func(data []byte, final bool) error + // Note: if we remove support for batching, this could just be a bool. + pendingJSONMessages []json.RawMessage - // streamRequests is the set of unanswered incoming requests for the stream. + // isBatch reports whether the request that opened the stream was a JSON-RPC + // batch. Batch responses must remain arrays even when they contain one item. + isBatch bool + + // w is the HTTP response writer for this stream. A non-nil w indicates + // that the stream is claimed by an HTTP request (the hanging POST or GET); + // it is set to nil when the request completes. + w http.ResponseWriter + + // done is closed to release the hanging HTTP request. + // + // Invariant: a non-nil done implies w is also non-nil, though the converse + // is not necessarily true: done is set to nil when it is closed, to avoid + // duplicate closure. + done chan struct{} + + // lastIdx is the index of the last written SSE event, for event ID generation. + // It starts at -1 since indices start at 0. + lastIdx int + + // protocolVersion is the protocol version for this stream. + protocolVersion string + + // requests is the set of unanswered incoming requests for the stream. // // Requests are removed when their response has been received. + // In practice, there is only one request, but in the 2025-03-26 version of + // the spec and earlier there was a concept of batching, in which POST + // payloads could hold multiple requests or responses. requests map[jsonrpc.ID]struct{} + + // isListen reports whether this stream was opened by a + // subscriptions/listen request. Listen streams are always SSE, live for + // the duration of the subscription, and act as the target for + // out-of-band notifications routed through this connection. + isListen bool +} + +// close sends a 'close' event to the client (if protocolVersion >= 2025-11-25 +// and reconnectAfter > 0) and closes the done channel. +// +// The done channel is set to nil after closing, so that done != nil implies +// the stream is active and done is open. This simplifies checks elsewhere. +func (s *stream) close(reconnectAfter time.Duration) { + s.mu.Lock() + defer s.mu.Unlock() + if s.done == nil { + return // stream not connected or already closed + } + if s.protocolVersion >= protocolVersion20251125 && reconnectAfter > 0 { + reconnectStr := strconv.FormatInt(reconnectAfter.Milliseconds(), 10) + if _, err := writeEvent(s.w, Event{ + Name: "close", + Retry: reconnectStr, + }); err != nil { + s.logger.Warn(fmt.Sprintf("Writing close event: %v", err)) + } + } + close(s.done) + s.done = nil +} + +// release releases the stream from its HTTP request, allowing it to be +// claimed by another request (e.g., for resumption). +func (s *stream) release() { + s.mu.Lock() + defer s.mu.Unlock() + s.w = nil + s.done = nil // may already be nil, if the stream is done or closed +} + +// extractErrorStatus reports the HTTP status to send when the given +// outgoing message is a JSON-RPC error response under the SEP-2575 protocol +// (>= 2026-07-28). +// +// Per SEP-2575: +// - MethodNotFound (-32601) MUST return HTTP 404. +// - InvalidParams (-32602), UnsupportedProtocolVersion (-32022) and +// CodeMissingRequiredClientCapabilities (-32021) MUST +// return HTTP 400. +func extractErrorStatus(ctx context.Context, msg jsonrpc.Message) int { + if protocolVersionFromContext(ctx) < protocolVersion20260728 { + return 0 + } + resp, ok := msg.(*jsonrpc.Response) + if !ok || resp.Error == nil { + return 0 + } + var jerr *jsonrpc.Error + if !errors.As(resp.Error, &jerr) { + return 0 + } + switch jerr.Code { + case jsonrpc.CodeMethodNotFound: + return http.StatusNotFound + case jsonrpc.CodeInvalidParams, CodeUnsupportedProtocolVersion, CodeMissingRequiredClientCapabilities: + return http.StatusBadRequest + } + return 0 +} + +// deliverLocked writes data to the stream (for SSE) or stores it in +// pendingJSONMessages (for JSON mode). The eventID is used for SSE event ID; +// pass "" to omit. +// +// If responseTo is valid, it is removed from the requests map. When all +// requests have been responded to, the done channel is closed and set to nil. +// +// If overrideStatus is non-zero, data is treated as a SEP-2575 protocol-level +// error response (>= 2026-07-28): it is written as a single raw JSON-RPC +// response body with Content-Type: application/json and HTTP status +// overrideStatus. +// +// Returns true if the stream is now done (all requests have been responded to). +// The done value is always accurate, even if an error is returned. +// +// s.mu must be held when calling this method. +func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.ID, overrideStatus int) (done bool, err error) { + // First, record the response. We must do this *before* returning an error + // below, as even if the stream is disconnected we want to update our + // accounting. + if responseTo.IsValid() { + delete(s.requests, responseTo) + } + // Now, try to deliver the message to the client. + done = len(s.requests) == 0 && s.id != "" + if s.done == nil { + return done, fmt.Errorf("stream not connected or already closed") + } + if done { + defer func() { close(s.done); s.done = nil }() + } + // SEP-2575 protocol-level error override: write the error as a raw + // JSON-RPC response with the spec-mandated HTTP status, bypassing any + // SSE framing. + if overrideStatus != 0 { + s.w.Header().Set("Content-Type", "application/json") + s.w.WriteHeader(overrideStatus) + if _, err := s.w.Write(data); err != nil { + return done, err + } + return done, nil + } + // Try to write to the response. + // + // If we get here, the request is still hanging (because s.done != nil + // implies s.w != nil), but may have been cancelled by the client/http layer: + // there's a brief race between request cancellation and releasing the + // stream. + if s.pendingJSONMessages != nil { + s.pendingJSONMessages = append(s.pendingJSONMessages, data) + if done { + // Flush all pending messages as JSON response. + var toWrite []byte + if len(s.pendingJSONMessages) == 1 && !s.isBatch { + toWrite = s.pendingJSONMessages[0] + } else { + toWrite, err = json.Marshal(s.pendingJSONMessages) + if err != nil { + return done, err + } + } + if _, err := s.w.Write(toWrite); err != nil { + return done, err + } + } + } else { + // SSE mode: write event to response writer. + s.lastIdx++ + if _, err := writeEvent(s.w, Event{Name: "message", Data: data, ID: eventID}); err != nil { + return done, err + } + } + return done, nil } // doneLocked reports whether the stream is logically complete. // +// s.requests was populated when reading the POST body, requests are deleted as +// they are responded to. Once all requests have been responded to, the stream +// is done. +// // s.mu must be held while calling this function. func (s *stream) doneLocked() bool { return len(s.requests) == 0 && s.id != "" } func (c *streamableServerConn) newStream(ctx context.Context, requests map[jsonrpc.ID]struct{}, id string) (*stream, error) { - if c.eventStore != nil { + if c.eventStore != nil && protocolVersionFromContext(ctx) < protocolVersion20260728 { if err := c.eventStore.Open(ctx, c.sessionID, id); err != nil { return nil, err } @@ -614,6 +1159,8 @@ func (c *streamableServerConn) newStream(ctx context.Context, requests map[jsonr return &stream{ id: id, requests: requests, + lastIdx: -1, // indices start at 0, incremented before each write + logger: c.logger, }, nil } @@ -640,6 +1187,18 @@ func (c *streamableServerConn) newStream(ctx context.Context, requests map[jsonr // ID, we avoid having to make this API decision. type idContextKey struct{} +// protocolVersionContextKey stores the protocol version extracted from the +// MCP-Protocol-Version HTTP header for use by lower layers. +type protocolVersionContextKey struct{} + +// protocolVersionFromContext returns the protocol version from the context, or +// the empty string if not set. An empty string means the version is unknown +// (e.g. the header was absent). +func protocolVersionFromContext(ctx context.Context) string { + v, _ := ctx.Value(protocolVersionContextKey{}).(string) + return v +} + // ServeHTTP handles a single HTTP request for the session. func (t *StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request) { if t.connection == nil { @@ -669,8 +1228,8 @@ func (c *streamableServerConn) serveGET(w http.ResponseWriter, req *http.Request // By default, we haven't seen a last index. Since indices start at 0, we represent // that by -1. This is incremented just before each event is written. lastIdx := -1 - if len(req.Header.Values("Last-Event-ID")) > 0 { - eid := req.Header.Get("Last-Event-ID") + if len(req.Header.Values(lastEventIDHeader)) > 0 { + eid := req.Header.Get(lastEventIDHeader) var ok bool streamID, lastIdx, ok = parseEventID(eid) if !ok { @@ -683,52 +1242,40 @@ func (c *streamableServerConn) serveGET(w http.ResponseWriter, req *http.Request } } - ctx, cancel := context.WithCancel(req.Context()) - defer cancel() + ctx := req.Context() - stream, done := c.acquireStream(ctx, w, streamID, &lastIdx) + protocolVersion := protocolVersionFromContext(ctx) + if protocolVersion == "" { + protocolVersion = protocolVersion20250326 + } + + stream, done := c.acquireStream(ctx, w, streamID, lastIdx, protocolVersion) if stream == nil { return } - // Release the stream when we're done. - defer func() { - stream.mu.Lock() - stream.deliver = nil - stream.mu.Unlock() - }() + defer stream.release() + c.hangResponse(ctx, done) +} +// hangResponse blocks the HTTP response until one of three conditions is met: +// - ctx is cancelled (the client disconnected or the request timed out) +// - done is closed (all responses have been sent, or the stream was explicitly closed) +// - the session is closed +// +// This keeps the HTTP connection open so that server-sent events can be +// written to the response. +func (c *streamableServerConn) hangResponse(ctx context.Context, done <-chan struct{}) { select { case <-ctx.Done(): - // request cancelled case <-done: - // request complete case <-c.done: - // session closed } } -// writeEvent writes an SSE event to w corresponding to the given stream, data, and index. -// lastIdx is incremented before writing, so that it continues to point to the index of the -// last event written to the stream. -func (c *streamableServerConn) writeEvent(w http.ResponseWriter, stream *stream, data []byte, lastIdx *int) error { - *lastIdx++ - e := Event{ - Name: "message", - Data: data, - } - if c.eventStore != nil { - e.ID = formatEventID(stream.id, *lastIdx) - } - if _, err := writeEvent(w, e); err != nil { - return err - } - return nil -} - -// acquireStream acquires the stream and replays all events since lastIdx, if -// any, updating lastIdx accordingly. If non-nil, the resulting stream will be -// registered for receiving new messages, and the resulting done channel will -// be closed when all related messages have been delivered. +// acquireStream replays all events since lastIdx, and acquires the ongoing +// stream, if any. If non-nil, the resulting stream will be registered for +// receiving new messages, and the stream's done channel will be closed when +// all related messages have been delivered. // // If any errors occur, they will be written to w and the resulting stream will // be nil. The resulting stream may also be nil if the stream is complete. @@ -736,10 +1283,15 @@ func (c *streamableServerConn) writeEvent(w http.ResponseWriter, stream *stream, // Importantly, this function must hold the stream mutex until done replaying // all messages, so that no delivery or storage of new messages occurs while // the stream is still replaying. -func (c *streamableServerConn) acquireStream(ctx context.Context, w http.ResponseWriter, streamID string, lastIdx *int) (*stream, chan struct{}) { +// +// protocolVersion is the protocol version for this stream, used to determine +// feature support (e.g. prime and close events were added in 2025-11-25). +func (c *streamableServerConn) acquireStream(ctx context.Context, w http.ResponseWriter, streamID string, lastIdx int, protocolVersion string) (*stream, chan struct{}) { // if tempStream is set, the stream is done and we're just replaying messages. // - // We record a temporary stream to claim exclusive replay rights. + // We record a temporary stream to claim exclusive replay rights. The spec + // (https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#resumability-and-redelivery) + // does not explicitly require exclusive replay, but we enforce it defensively. tempStream := false c.mu.Lock() s, ok := c.streams[streamID] @@ -747,12 +1299,12 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons // The stream is logically done, but claim exclusive rights to replay it by // adding a temporary entry in the streams map. // - // We create this entry with a non-nil deliver function, to ensure it isn't - // claimed by another request before we lock it below. + // We create this entry with a non-nil w, to ensure it isn't claimed by + // another request before we lock it below. tempStream = true s = &stream{ - id: streamID, - deliver: func([]byte, bool) error { return nil }, + id: streamID, + w: w, } c.streams[streamID] = s @@ -769,7 +1321,7 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons defer s.mu.Unlock() // Check that this stream wasn't claimed by another request. - if !tempStream && s.deliver != nil { + if !tempStream && s.w != nil { http.Error(w, "stream ID conflicts with ongoing stream", http.StatusConflict) return nil, nil } @@ -782,7 +1334,7 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons // messages, and registered our delivery function. var toReplay [][]byte if c.eventStore != nil { - for data, err := range c.eventStore.After(ctx, c.SessionID(), s.id, *lastIdx) { + for data, err := range c.eventStore.After(ctx, c.SessionID(), s.id, lastIdx) { if err != nil { // We can't replay events, perhaps because the underlying event store // has garbage collected its storage. @@ -795,7 +1347,9 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons http.Error(w, "failed to replay events", http.StatusBadRequest) return nil, nil } - toReplay = append(toReplay, data) + if len(data) > 0 { + toReplay = append(toReplay, data) + } } } @@ -806,14 +1360,35 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons if s.id == "" { // Issue #410: the standalone SSE stream is likely not to receive messages // for a long time. Ensure that headers are flushed. + // + // On HTTP/2, headers and body travel as separate frames (HEADERS and + // DATA). Reverse proxies (e.g. Envoy, Caddy, net/http/httputil) + // commonly buffer the HEADERS frame until they have a DATA frame to + // coalesce it with — there is no HTTP/2 equivalent of HTTP/1.1's + // Transfer-Encoding: chunked signal that says "this is streaming, send + // headers now". Calling Flush() alone is not sufficient: it pushes + // the kernel buffer to the proxy, but the proxy still holds the + // HEADERS frame. + // + // Write an SSE comment (lines starting with ":" are ignored by + // clients per RFC) so a DATA frame is produced, which forces the + // proxy to forward both frames. See: + // https://github.com/golang/go/issues/31125 + // https://github.com/caddyserver/caddy/issues/4247 w.WriteHeader(http.StatusOK) - if f, ok := w.(http.Flusher); ok { - f.Flush() - } + fmt.Fprint(w, ": ok\n\n") + rc := http.NewResponseController(w) + // Ignore returned error as flushing is best-effort. + _ = rc.Flush() } for _, data := range toReplay { - if err := c.writeEvent(w, s, data, lastIdx); err != nil { + lastIdx++ + e := Event{Name: "message", Data: data} + if c.eventStore != nil { + e.ID = formatEventID(s.id, lastIdx) + } + if _, err := writeEvent(w, e); err != nil { return nil, nil } } @@ -823,20 +1398,13 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons return nil, nil } - // The stream is not done: register a delivery function before the stream is + // The stream is not done: set up delivery state before the stream is // unlocked, allowing the connection to write new events. - done := make(chan struct{}) - s.deliver = func(data []byte, final bool) error { - if err := ctx.Err(); err != nil { - return err - } - err := c.writeEvent(w, s, data, lastIdx) - if final { - close(done) - } - return err - } - return s, done + s.w = w + s.done = make(chan struct{}) + s.lastIdx = lastIdx + s.protocolVersion = protocolVersion + return s, s.done } // servePOST handles an incoming message, and replies with either an outgoing @@ -845,7 +1413,7 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons // // It returns an HTTP status code and error message. func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Request) { - if len(req.Header.Values("Last-Event-ID")) > 0 { + if len(req.Header.Values(lastEventIDHeader)) > 0 { http.Error(w, "can't send Last-Event-ID for POST request", http.StatusBadRequest) return } @@ -853,6 +1421,11 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques // Read incoming messages. body, err := io.ReadAll(req.Body) if err != nil { + var mbe *http.MaxBytesError + if errors.As(err, &mbe) { + http.Error(w, fmt.Sprintf("request body exceeds %d bytes", mbe.Limit), http.StatusRequestEntityTooLarge) + return + } http.Error(w, "failed to read body", http.StatusBadRequest) return } @@ -860,13 +1433,16 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques http.Error(w, "POST requires a non-empty body", http.StatusBadRequest) return } + // TODO(#674): once we've documented the support matrix for 2025-03-26 and + // earlier, drop support for matching entirely; that will simplify this + // logic. incoming, isBatch, err := readBatch(body) if err != nil { http.Error(w, fmt.Sprintf("malformed payload: %v", err), http.StatusBadRequest) return } - protocolVersion := req.Header.Get(protocolVersionHeader) + protocolVersion := protocolVersionFromContext(req.Context()) if protocolVersion == "" { protocolVersion = protocolVersion20250326 } @@ -886,30 +1462,190 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques calls := make(map[jsonrpc.ID]struct{}) tokenInfo := auth.TokenInfoFromContext(req.Context()) isInitialize := false + isSubscriptionsListen := false + var initializeProtocolVersion string for _, msg := range incoming { if jreq, ok := msg.(*jsonrpc.Request); ok { // Preemptively check that this is a valid request, so that we can fail // the HTTP request. If we didn't do this, a request with a bad method or // missing ID could be silently swallowed. - if _, err := checkRequest(jreq, serverMethodInfos); err != nil { + // Use the server's receiving method infos (which include any custom + // methods registered via AddReceivingCustomMethod) when available; + // fall back to the standard methods otherwise, e.g. in tests that + // exercise streamableServerConn directly without a server. + methodInfos := serverMethodInfos + if c.server != nil { + methodInfos = c.server.receivingMethodInfos() + } + if _, err := checkRequest(jreq, methodInfos); err != nil { + if protocolVersion >= protocolVersion20260728 && errors.Is(err, jsonrpc2.ErrNotHandled) && jreq.IsCall() { + writeJSONRPCError(w, http.StatusNotFound, jreq.ID, &jsonrpc.Error{ + Code: jsonrpc.CodeMethodNotFound, + Message: err.Error(), + }) + return + } http.Error(w, err.Error(), http.StatusBadRequest) return } if jreq.Method == methodInitialize { isInitialize = true + // Extract the protocol version from InitializeParams. + var params InitializeParams + if err := internaljson.Unmarshal(jreq.Params, ¶ms); err == nil { + initializeProtocolVersion = params.ProtocolVersion + } } + if jreq.Method == methodSubscriptionsListen { + isSubscriptionsListen = true + } + // SEP-2575: requests carrying `_meta.protocolVersion` require the + // Mcp-Protocol-Version HTTP header to be present and to match the + // per-request `_meta.protocolVersion` value. + // The new (>= 2026-07-28) protocol is supported on the HTTP transport + // only when [StreamableHTTPOptions.Stateless] is true. + // + // The `_meta` triple is defined for calls only, notifications are + // excluded. + // + // TODO: this validation can be moved within validateMcpHeaders. + var metaVersion string + if meta := extractRequestMeta(jreq.Params); meta != nil { + metaVersion, _ = meta[MetaKeyProtocolVersion].(string) + } + if (jreq.IsCall() && protocolVersion >= protocolVersion20260728) || metaVersion != "" { + // Extract again the protocol version from the context to see what the client + // is advertising in the Mcp-Protocol-Version HTTP header. + headerVersion := protocolVersionFromContext(req.Context()) + // server/discover is exempt from the stateful + // rejection as it should learn about the supported protocols from the + // DiscoverResult response. + if !c.stateless && jreq.Method != methodDiscover { + if plaintextstatefulrejection == "1" { + http.Error(w, fmt.Sprintf( + "Bad Request: protocol version %q is only supported on stateless HTTP servers (set StreamableHTTPOptions.Stateless = true)", + protocolVersion), + http.StatusBadRequest) + return + } + // Advertise only the legacy versions this stateful server accepts + // (i.e. exclude 2026-07-28, since that's what the request asked for + // and what we're rejecting). + advertised := supportedProtocolVersions + if c.server != nil { + advertised = c.server.protocolVersions + } + legacyVersions := slices.DeleteFunc(slices.Clone(advertised), func(v string) bool { + return v >= protocolVersion20260728 + }) + data, _ := json.Marshal(UnsupportedProtocolVersionData{ + Supported: legacyVersions, + Requested: protocolVersion, + }) + c.logger.Warn(fmt.Sprintf( + "rejecting request with protocol version %q: this server is stateful; set StreamableHTTPOptions.Stateless = true to accept it", + protocolVersion)) + writeJSONRPCError(w, http.StatusBadRequest, jreq.ID, &jsonrpc.Error{ + Code: CodeUnsupportedProtocolVersion, + Message: fmt.Sprintf("protocol version %q is not supported by this server", protocolVersion), + Data: data, + }) + return + } + if headerVersion == "" { + writeJSONRPCError(w, http.StatusBadRequest, jreq.ID, &jsonrpc.Error{ + Code: CodeHeaderMismatch, + Message: fmt.Sprintf( + "%s header is required for requests carrying %q", + protocolVersionHeader, MetaKeyProtocolVersion), + }) + return + } + if metaVersion == "" { + writeJSONRPCError(w, http.StatusBadRequest, jreq.ID, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: fmt.Sprintf( + "missing or invalid _meta field %q", + MetaKeyProtocolVersion), + }) + return + } + if headerVersion != metaVersion { + writeJSONRPCError(w, http.StatusBadRequest, jreq.ID, &jsonrpc.Error{ + Code: CodeHeaderMismatch, + Message: fmt.Sprintf( + "%s header %q does not match request %s %q", + protocolVersionHeader, headerVersion, + MetaKeyProtocolVersion, metaVersion), + }) + return + } + } + // Include metadata for all requests (including notifications). jreq.Extra = &RequestExtra{ TokenInfo: tokenInfo, Header: req.Header, } if jreq.IsCall() { calls[jreq.ID] = struct{}{} + // See the doc for CloseSSEStream: allow the request handler to + // explicitly close the ongoing stream. + jreq.Extra.(*RequestExtra).CloseSSEStream = func(args CloseSSEStreamArgs) { + // This mechanism was designed to trigger client reconnection with + // Last-Event-ID for server-initiated disconnect scenarios. It is + // deprecated in protocol version 2026-07-28. + if protocolVersion >= protocolVersion20260728 { + return + } + c.mu.Lock() + streamID, ok := c.requestStreams[jreq.ID] + var stream *stream + if ok { + stream = c.streams[streamID] + } + c.mu.Unlock() + + if stream != nil { + stream.close(args.RetryAfter) + } + } } } } + // Validate MCP standard headers (Mcp-Method, Mcp-Name, Mcp-Param-*) + if !isBatch && len(incoming) == 1 { + if err := validateMcpHeaders(req.Header, incoming[0], c.toolLookup); err != nil { + resp := &jsonrpc.Response{ + Error: jsonrpc2.NewError(CodeHeaderMismatch, err.Error()), + } + if jreq, ok := incoming[0].(*jsonrpc.Request); ok { + resp.ID = jreq.ID + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + if data, err := jsonrpc2.EncodeMessage(resp); err == nil { + w.Write(data) + } + return + } + } + + // The prime and close events were added in protocol version 2025-11-25 (SEP-1699). + // Use the version from InitializeParams if this is an initialize request, + // otherwise use the protocol version header. + effectiveVersion := protocolVersion + if isInitialize && initializeProtocolVersion != "" { + effectiveVersion = initializeProtocolVersion + } + // If we don't have any calls, we can just publish the incoming messages and return. // No need to track a logical stream. + // + // See section [§2.1.4] of the spec: "If the server accepts the input, the + // server MUST return HTTP status code 202 Accepted with no body." + // + // [§2.1.4]: https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#sending-messages-to-the-server if len(calls) == 0 { for _, msg := range incoming { select { @@ -929,87 +1665,92 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques // // Create a logical stream to track its responses. // Important: don't publish the incoming messages until the stream is - // registered, as the server may attempt to respond to imcoming messages as + // registered, as the server may attempt to respond to incoming messages as // soon as they're published. - stream, err := c.newStream(req.Context(), calls, randText()) + stream, err := c.newStream(req.Context(), calls, crand.Text()) if err != nil { http.Error(w, fmt.Sprintf("storing stream: %v", err), http.StatusInternalServerError) return } + stream.isListen = isSubscriptionsListen + stream.isBatch = isBatch + + // subscriptions/listen is inherently a long-lived SSE endpoint (SEP-2575): + // it has no synchronous result, the response stream stays open until the + // client cancels, and the server pushes notifications on it as they occur. + // Force SSE mode (bypassing JSONResponse) so the buffered application/json + // path doesn't deadlock waiting for a completion that won't come. + useSSE := !c.jsonResponse || isSubscriptionsListen // Set response headers. Accept was checked in [StreamableHTTPHandler]. w.Header().Set("Cache-Control", "no-cache, no-transform") - if c.jsonResponse { - w.Header().Set("Content-Type", "application/json") - } else { + if useSSE { w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Connection", "keep-alive") + } else { + w.Header().Set("Content-Type", "application/json") } if c.sessionID != "" && isInitialize { w.Header().Set(sessionIDHeader, c.sessionID) } - // Message delivery has two paths, depending on whether we're responding with JSON or - // event stream. - done := make(chan struct{}) // closed after the final response is written - if c.jsonResponse { - var msgs []json.RawMessage - stream.deliver = func(data []byte, final bool) error { - // Collect messages until we've received the final response. - // - // In recent protocol versions, there should only be one message as - // batching is disabled, as checked above. - msgs = append(msgs, data) - if !final { - return nil - } - defer close(done) // final response + // Set up stream delivery state. + stream.w = w + done := make(chan struct{}) + stream.done = done + stream.protocolVersion = effectiveVersion - // Write either the JSON object corresponding to the one response, or a - // JSON array corresponding to the batch response. - var toWrite []byte - if len(msgs) == 1 { - toWrite = []byte(msgs[0]) - } else { - var err error - toWrite, err = json.Marshal(msgs) - if err != nil { - return err - } - } - _, err = w.Write(toWrite) - return err - } - } else { - // Write events in the order we receive them. - lastIndex := -1 - stream.deliver = func(data []byte, final bool) error { - if final { - defer close(done) - } - return c.writeEvent(w, stream, data, &lastIndex) + // Reject any call whose ID is already in flight on this session, + // atomically and without partial registration. + c.mu.Lock() + for reqID := range calls { + if _, ok := c.requestStreams[reqID]; ok { + c.mu.Unlock() + writeJSONRPCError(w, http.StatusBadRequest, reqID, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidRequest, + Message: fmt.Sprintf("duplicate in-flight request ID %v", reqID.Raw()), + }) + return } } - - // Release ownership of the stream by unsetting deliver. - defer func() { - stream.mu.Lock() - // TODO(rfindley): if we have no event store, we should really cancel all - // remaining requests here, since the client will never get the results. - stream.deliver = nil - stream.mu.Unlock() - }() - - // The stream is now set up to deliver messages. - // - // Register it before publishing incoming messages. - c.mu.Lock() c.streams[stream.id] = stream for reqID := range calls { c.requestStreams[reqID] = stream.id } c.mu.Unlock() + // TODO(rfindley): if we have no event store, we should really cancel all + // remaining requests here, since the client will never get the results. + defer stream.release() + + if !useSSE { + // JSON mode: collect messages in pendingJSONMessages until done. + // Set pendingJSONMessages to a non-nil value to signal that this is an + // application/json stream. + stream.pendingJSONMessages = []json.RawMessage{} + } else { + // SSE mode: write a priming event if supported. + // + // SEP-2575 removes Last-Event-ID-based resumable streams for protocol + // version >= 2026-07-28. + if c.eventStore != nil && effectiveVersion >= protocolVersion20251125 && effectiveVersion < protocolVersion20260728 { + // Write a priming event, as defined by [§2.1.6] of the spec. + // + // [§2.1.6]: https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#sending-messages-to-the-server + // + // We must also write it to the event store in order for indexes to + // align. + if err := c.eventStore.Append(req.Context(), c.sessionID, stream.id, nil); err != nil { + c.logger.Warn(fmt.Sprintf("Storing priming event: %v", err)) + } + stream.lastIdx++ + e := Event{Name: "prime", ID: formatEventID(stream.id, stream.lastIdx)} + if _, err := writeEvent(w, e); err != nil { + c.logger.Warn(fmt.Sprintf("Writing priming event: %v", err)) + } + } + } + // Publish incoming messages. for _, msg := range incoming { select { @@ -1017,6 +1758,8 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques // Note: don't select on req.Context().Done() here, since we've already // received the requests and may have already published a response message // or notification. The client could resume the stream. + // + // In fact, this send could be in a separate goroutine. case <-c.done: // Session closed: we don't know if any data has been written, so it's // too late to write a status code here. @@ -1024,14 +1767,7 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques } } - select { - case <-req.Context().Done(): - // request cancelled - case <-done: - // request complete - case <-c.done: - // session is closed - } + c.hangResponse(req.Context(), done) } // Event IDs: encode both the logical connection ID and the index, as @@ -1090,7 +1826,7 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e return err } - if req, ok := msg.(*jsonrpc.Request); ok && req.ID.IsValid() && (c.stateless || c.sessionID == "") { + if req, ok := msg.(*jsonrpc.Request); ok && req.IsCall() && (c.stateless || c.sessionID == "") { // Requests aren't possible with stateless servers, or when there's no session ID. return fmt.Errorf("%w: stateless servers cannot make requests", jsonrpc2.ErrRejected) } @@ -1127,7 +1863,18 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e s = c.streams[streamID] } } else { - s = c.streams[""] // standalone SSE stream + // In stateless mode there will always be only one stream per connection. + // If that stream was open to listen for subscription notifications, + // automatically select as the one to write the notification to. + for _, stream := range c.streams { + if stream.isListen { + s = stream + break + } + } + if s == nil { + s = c.streams[""] // standalone SSE stream + } } if responseTo.IsValid() { // Once we've responded to a request, disallow related messages by removing @@ -1151,48 +1898,48 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e s.mu.Lock() defer s.mu.Unlock() - if s.doneLocked() { - // It's possible that the stream was completed in between getting s above, - // and acquiring the stream lock. In order to avoid acquiring s.mu while - // holding c.mu, we check the terminal condition again. - return fmt.Errorf("%w: write to closed stream", jsonrpc2.ErrRejected) - } - // Perform accounting on responses. - if responseTo.IsValid() { - if _, ok := s.requests[responseTo]; !ok { - panic(fmt.Sprintf("internal error: stream %v: response to untracked request %v", s.id, responseTo)) - } - if s.id == "" { - // This should be guaranteed not to happen by the stream resolution logic - // above, but be defensive: we don't ever want to delete the standalone - // stream. - panic("internal error: response on standalone stream") - } - delete(s.requests, responseTo) - if len(s.requests) == 0 { - c.mu.Lock() - delete(c.streams, s.id) - c.mu.Unlock() + + // Store in eventStore before delivering. + // TODO(rfindley): we should only append if the response is SSE, not JSON, by + // pushing down into the delivery layer. + delivered := false + var errs []error + protocolVersion := protocolVersionFromContext(ctx) + if c.eventStore != nil && protocolVersion < protocolVersion20260728 { + if err := c.eventStore.Append(ctx, c.sessionID, s.id, data); err != nil { + errs = append(errs, err) + } else { + delivered = true } } - delivered := false - if c.eventStore != nil { - if err := c.eventStore.Append(ctx, c.sessionID, s.id, data); err != nil { - // TODO: report a side-channel error. - } else { - delivered = true - } + // Compute eventID for SSE streams with event store. + // Use s.lastIdx + 1 because deliverLocked increments before writing. + var eventID string + if c.eventStore != nil && protocolVersion < protocolVersion20260728 { + eventID = formatEventID(s.id, s.lastIdx+1) } - if s.deliver != nil { - if err := s.deliver(data, s.doneLocked()); err != nil { - // TODO: report a side-channel error. - } else { - delivered = true - } + + // SEP-2575: map protocol-level JSON-RPC error codes to HTTP status codes + // on the new protocol (>= 2026-07-28). When non-zero, deliverLocked will + // write the body as raw application/json with the override status. + overrideStatus := extractErrorStatus(ctx, msg) + + done, err := s.deliverLocked(data, eventID, responseTo, overrideStatus) + if err != nil { + errs = append(errs, err) + } else { + delivered = true } + + if done { + c.mu.Lock() + delete(c.streams, s.id) + c.mu.Unlock() + } + if !delivered { - return fmt.Errorf("%w: undelivered message", jsonrpc2.ErrRejected) + return fmt.Errorf("%w: undelivered message: %v", jsonrpc2.ErrRejected, errors.Join(errs...)) } return nil } @@ -1223,6 +1970,34 @@ type StreamableClientTransport struct { // It defaults to 5. To disable retries, use a negative number. MaxRetries int + // DisableStandaloneSSE controls whether the client establishes a standalone SSE stream + // for receiving server-initiated messages. + // + // When false (the default), after initialization the client sends an HTTP GET request + // to establish a persistent server-sent events (SSE) connection. This allows the server + // to send messages to the client at any time, such as ToolListChangedNotification or + // other server-initiated requests and notifications. The connection persists for the + // lifetime of the session and automatically reconnects if interrupted. + // + // When true, the client does not establish the standalone SSE stream. The client will + // only receive responses to its own POST requests. Server-initiated messages will not + // be received. + // + // According to the MCP specification, the standalone SSE stream is optional. + // Setting DisableStandaloneSSE to true is useful when: + // - You only need request-response communication and don't need server-initiated notifications + // - The server doesn't properly handle GET requests for SSE streams + // - You want to avoid maintaining a persistent connection + DisableStandaloneSSE bool + + // OAuthHandler is an optional field that, if provided, will be used to authorize the requests. + OAuthHandler auth.OAuthHandler + + // MaxEventSize bounds the number of bytes buffered while reading a single + // server-sent event. A value of 0 selects [DefaultMaxEventSize], a negative + // value disables the cap. + MaxEventSize int + // TODO(rfindley): propose exporting these. // If strict is set, the transport is in 'strict mode', where any violation // of the MCP spec causes a failure. @@ -1239,20 +2014,36 @@ const ( // A value of 1.0 results in a constant delay, while a value of 2.0 would double it each time. // It must be 1.0 or greater if MaxRetries is greater than 0. reconnectGrowFactor = 1.5 - // reconnectInitialDelay is the base delay for the first reconnect attempt. - reconnectInitialDelay = 1 * time.Second // reconnectMaxDelay caps the backoff delay, preventing it from growing indefinitely. reconnectMaxDelay = 30 * time.Second + // closeDeleteTimeout bounds the session-termination DELETE sent by Close. + // That request is best-effort: a peer that accepts the connection but never + // answers must not be able to hold teardown open until the TCP timeout. + // See issue #1183. + closeDeleteTimeout = 5 * time.Second ) +var ( + // reconnectInitialDelay is the base delay for the first reconnect attempt. + // + // Mutable for testing. + reconnectInitialDelay atomic.Int64 +) + +func init() { + reconnectInitialDelay.Store(int64(1 * time.Second)) +} + // Connect implements the [Transport] interface. // // The resulting [Connection] writes messages via POST requests to the -// transport URL with the Mcp-Session-Id header set, and reads messages from -// hanging requests. +// transport URL, and reads messages from hanging requests. If the server +// provides a session ID via the Mcp-Session-Id response header, subsequent +// requests include it; sessionless servers that omit the header are fully +// supported. // -// When closed, the connection issues a DELETE request to terminate the logical -// session. +// When closed, the connection issues a DELETE request to terminate the +// session, unless no session was established. func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, error) { client := t.HTTPClient if client == nil { @@ -1267,18 +2058,34 @@ func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, er // Create a new cancellable context that will manage the connection's lifecycle. // This is crucial for cleanly shutting down the background SSE listener by // cancelling its blocking network operations, which prevents hangs on exit. - connCtx, cancel := context.WithCancel(ctx) + // + // This context should be detached from the incoming context: the standalone + // SSE request should not break when the connection context is done. + // + // For example, consider that the user may want to wait at most 5s to connect + // to the server, and therefore uses a context with a 5s timeout when calling + // client.Connect. Let's suppose that Connect returns after 1s, and the user + // starts using the resulting session. If we didn't detach here, the session + // would break after 4s, when the background SSE stream is terminated. + // + // Instead, creating a cancellable context detached from the incoming context + // allows us to preserve context values (which may be necessary for auth + // middleware), yet only cancel the standalone stream when the connection is closed. + connCtx, cancel := context.WithCancel(xcontext.Detach(ctx)) conn := &streamableClientConn{ - url: t.Endpoint, - client: client, - incoming: make(chan jsonrpc.Message, 10), - done: make(chan struct{}), - maxRetries: maxRetries, - strict: t.strict, - logger: t.logger, - ctx: connCtx, - cancel: cancel, - failed: make(chan struct{}), + url: t.Endpoint, + client: client, + incoming: make(chan jsonrpc.Message, 10), + done: make(chan struct{}), + maxRetries: maxRetries, + strict: t.strict, + logger: ensureLogger(t.logger), // must be non-nil for safe logging + ctx: connCtx, + cancel: cancel, + failed: make(chan struct{}), + disableStandaloneSSE: t.DisableStandaloneSSE, + oauthHandler: t.OAuthHandler, + maxEventSize: t.MaxEventSize, } return conn, nil } @@ -1286,13 +2093,24 @@ func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, er type streamableClientConn struct { url string client *http.Client - ctx context.Context - cancel context.CancelFunc + ctx context.Context // connection context, detached from Connect + cancel context.CancelFunc // cancels ctx incoming chan jsonrpc.Message maxRetries int strict bool // from [StreamableClientTransport.strict] logger *slog.Logger // from [StreamableClientTransport.logger] + // disableStandaloneSSE controls whether to disable the standalone SSE stream + // for receiving server-to-client notifications when no request is in flight. + disableStandaloneSSE bool // from [StreamableClientTransport.DisableStandaloneSSE] + + // oauthHandler is the OAuth handler for the connection. + oauthHandler auth.OAuthHandler // from [StreamableClientTransport.OAuthHandler] + + // maxEventSize bounds the number of bytes buffered while reading a single + // server-sent event. Resolved from [StreamableClientTransport.MaxEventSize]. + maxEventSize int + // Guard calls to Close, as it may be called multiple times. closeOnce sync.Once closeErr error @@ -1315,17 +2133,6 @@ type streamableClientConn struct { sessionID string } -// errSessionMissing distinguishes if the session is known to not be present on -// the server (see [streamableClientConn.fail]). -// -// TODO(rfindley): should we expose this error value (and its corresponding -// API) to the user? -// -// The spec says that if the server returns 404, clients should reestablish -// a session. For now, we delegate that to the user, but do they need a way to -// differentiate a 'NotFound' error from other errors? -var errSessionMissing = errors.New("session not found") - var _ clientConnection = (*streamableClientConn)(nil) func (c *streamableClientConn) sessionUpdated(state clientSessionState) { @@ -1333,8 +2140,15 @@ func (c *streamableClientConn) sessionUpdated(state clientSessionState) { c.initializedResult = state.InitializeResult c.mu.Unlock() + // Under SEP-2575 (protocol version >= 2026-07-28) the standalone HTTP GET + // SSE stream is removed. + if state.InitializeResult == nil || + state.InitializeResult.ProtocolVersion >= protocolVersion20260728 { + return + } + // Start the standalone SSE stream as soon as we have the initialized - // result. + // result, if continuous listening is enabled. // // § 2.2: The client MAY issue an HTTP GET to the MCP endpoint. This can be // used to open an SSE stream, allowing the server to communicate to the @@ -1344,9 +2158,59 @@ func (c *streamableClientConn) sessionUpdated(state clientSessionState) { // initialized, we don't know whether the server requires a sessionID. // // § 2.5: A server using the Streamable HTTP transport MAY assign a session - // ID at initialization time, by including it in an Mcp-Session-Id header + // ID at initialization time, by including it in a Mcp-Session-Id header // on the HTTP response containing the InitializeResult. - go c.handleSSE("standalone SSE stream", nil, true, nil) + if !c.disableStandaloneSSE { + c.connectStandaloneSSE() + } +} + +func (c *streamableClientConn) connectStandaloneSSE() { + resp, err := c.connectSSE(c.ctx, "", 0, true) + if err != nil { + // If the client didn't cancel the request, and failure breaks the logical + // session. + if c.ctx.Err() == nil { + c.fail(fmt.Errorf("standalone SSE request failed (session ID: %v): %v", c.sessionID, err)) + } + return + } + + // [§2.2.3]: "The server MUST either return Content-Type: + // text/event-stream in response to this HTTP GET, or else return HTTP + // 405 Method Not Allowed, indicating that the server does not offer an + // SSE stream at this endpoint." + // + // [§2.2.3]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#listening-for-messages-from-the-server + if resp.StatusCode == http.StatusMethodNotAllowed { + // The server doesn't support the standalone SSE stream. + resp.Body.Close() + return + } + if baseMediaType(resp.Header.Get("Content-Type")) != "text/event-stream" { + // modelcontextprotocol/go-sdk#736: some servers return 200 OK or redirect with + // non-SSE content type instead of text/event-stream for the standalone + // SSE stream. + c.logger.Warn(fmt.Sprintf("got Content-Type %s instead of text/event-stream for standalone SSE stream", resp.Header.Get("Content-Type"))) + resp.Body.Close() + return + } + if resp.StatusCode >= 400 && resp.StatusCode < 500 && !c.strict { + // modelcontextprotocol/go-sdk#393,#610: some servers return NotFound or + // other status codes instead of MethodNotAllowed for the standalone SSE + // stream. + // + // Treat this like MethodNotAllowed in non-strict mode. + c.logger.Warn(fmt.Sprintf("got %d instead of 405 for standalone SSE stream", resp.StatusCode)) + resp.Body.Close() + return + } + summary := "standalone SSE stream" + if err := c.checkResponse(c.ctx, summary, resp); err != nil { + c.fail(err) + return + } + go c.handleSSE(c.ctx, summary, resp, nil) } // fail handles an asynchronous error while reading. @@ -1354,7 +2218,7 @@ func (c *streamableClientConn) sessionUpdated(state clientSessionState) { // If err is non-nil, it is terminal, and subsequent (or pending) Reads will // fail. // -// If err wraps errSessionMissing, the failure indicates that the session is no +// If err wraps ErrSessionMissing, the failure indicates that the session is no // longer present on the server, and no final DELETE will be performed when // closing the connection. func (c *streamableClientConn) fail(err error) { @@ -1405,11 +2269,15 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e } var requestSummary string - var isCall bool + var requestMethod string + var forCall *jsonrpc.Request switch msg := msg.(type) { case *jsonrpc.Request: requestSummary = fmt.Sprintf("sending %q", msg.Method) - isCall = msg.IsCall() + if msg.IsCall() { + forCall = msg + requestMethod = msg.Method + } case *jsonrpc.Response: requestSummary = fmt.Sprintf("sending jsonrpc response #%d", msg.ID) default: @@ -1421,35 +2289,85 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e return fmt.Errorf("%s: %v", requestSummary, err) } - req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(data)) + doRequest := func() (*http.Request, *http.Response, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(data)) + if err != nil { + return nil, nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + + if err := c.setMCPHeaders(req, msg); err != nil { + // Failure to set headers means that the request was not sent. + // Wrap with ErrRejected so the jsonrpc2 connection doesn't set writeErr + // and permanently break the connection. + return nil, nil, fmt.Errorf("%s: %w: %w", requestSummary, jsonrpc2.ErrRejected, err) + } + // Keep this after the setMCPHeaders call to ensure that the + // protocol version header is set. + setStandardHeaders(ctx, req.Header, msg) + resp, err := c.client.Do(req) + if err != nil { + // Any error from client.Do means the request didn't reach the server. + // Wrap with ErrRejected so the jsonrpc2 connection doesn't set writeErr + // and permanently break the connection. + err = fmt.Errorf("%s: %w: %w", requestSummary, jsonrpc2.ErrRejected, err) + } + return req, resp, err + } + + req, resp, err := doRequest() if err != nil { return err } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Accept", "application/json, text/event-stream") - c.setMCPHeaders(req) - resp, err := c.client.Do(req) - if err != nil { - return fmt.Errorf("%s: %v", requestSummary, err) + if (resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden) && c.oauthHandler != nil { + if err := c.oauthHandler.Authorize(ctx, req, resp); err != nil { + // If the caller's context was cancelled while we were running the + // authorization flow, treat the connection as failed so subsequent + // operations on it (e.g. the cancellation notify the call layer + // sends in response to ctx cancellation) short-circuit instead of + // re-invoking the OAuth handler. Otherwise the user gets prompted + // to authorize a request they have already abandoned. See #882. + // + // We check ctx.Err() rather than the error returned by Authorize, + // because the handler is user-implemented and may return an error + // that does not wrap context.Canceled (e.g. a custom sentinel or + // a fmt.Errorf with %v). The context itself is the authoritative + // source for whether the caller abandoned the request. + ctxErr := ctx.Err() + if errors.Is(ctxErr, context.Canceled) || errors.Is(ctxErr, context.DeadlineExceeded) { + c.fail(fmt.Errorf("%s: authorization cancelled: %w", requestSummary, err)) + } + // Wrap with ErrRejected so the jsonrpc2 connection doesn't set writeErr + // and permanently break the connection. + // Wrap the authorization error as well for client inspection. + return fmt.Errorf("%s: %w: %w", requestSummary, jsonrpc2.ErrRejected, err) + } + // Retry the request after successful authorization. + _, resp, err = doRequest() + if err != nil { + return err + } } - // §2.5.3: "The server MAY terminate the session at any time, after - // which it MUST respond to requests containing that session ID with HTTP - // 404 Not Found." - if resp.StatusCode == http.StatusNotFound { - // Fail the session immediately, rather than relying on jsonrpc2 to fail - // (and close) it, because we want the call to Close to know that this - // session is missing (and therefore not send the DELETE). - err := fmt.Errorf("%s: failed to send: %w", requestSummary, errSessionMissing) - c.fail(err) - resp.Body.Close() + if err := c.checkResponse(ctx, requestSummary, resp); err != nil { + if (requestMethod == methodDiscover || requestMethod == methodSubscriptionsListen) && !errors.Is(err, jsonrpc2.ErrRejected) { + // Wrap the discover or subscriptions/listen failure with ErrRejected so + // the jsonrpc2 layer doesn't set writeErr, which would break the connection + // (preventing the legacy initialize fallback from succeeding, or breaking + // subsequent RPCs if subscriptions/listen fails). This covers the case + // where a server rejects these requests with a non-JSON-RPC body (e.g. + // plain text 404 or 400), which checkResponse cannot classify as a + // per-call rejection on its own. + err = fmt.Errorf("%w: %w", err, jsonrpc2.ErrRejected) + } else if !errors.Is(err, jsonrpc2.ErrRejected) { + // Only fail the connection for non-transient errors. + // Transient errors (wrapped with ErrRejected) should not break the connection. + c.fail(err) + } return err } - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - resp.Body.Close() - return fmt.Errorf("broken session: %v", resp.Status) - } if sessionID := resp.Header.Get(sessionIDHeader); sessionID != "" { c.mu.Lock() @@ -1463,24 +2381,28 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e return fmt.Errorf("mismatching session IDs %q and %q", hadSessionID, sessionID) } } - if resp.StatusCode == http.StatusNoContent || resp.StatusCode == http.StatusAccepted { + + if forCall == nil { + resp.Body.Close() + // [§2.1.4]: "If the input is a JSON-RPC response or notification: // If the server accepts the input, the server MUST return HTTP status code 202 Accepted with no body." // // [§2.1.4]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#listening-for-messages-from-the-server - resp.Body.Close() - return nil - } else if !isCall && !c.strict { - // Some servers return 200, even with an empty json body. - // Ignore this response in non-strict mode. - if c.logger != nil { - c.logger.Warn(fmt.Sprintf("unexpected status code %d from non-call", resp.StatusCode)) + if resp.StatusCode != http.StatusNoContent && resp.StatusCode != http.StatusAccepted { + errMsg := fmt.Sprintf("unexpected status code %d from non-call", resp.StatusCode) + // Some servers return 200, even with an empty json body. + // + // In strict mode, return an error to the caller. + c.logger.Warn(errMsg) + if c.strict { + return errors.New(errMsg) + } } - resp.Body.Close() return nil } - contentType := strings.TrimSpace(strings.SplitN(resp.Header.Get("Content-Type"), ";", 2)[0]) + contentType := baseMediaType(resp.Header.Get("Content-Type")) switch contentType { case "application/json": go c.handleJSON(requestSummary, resp) @@ -1490,8 +2412,11 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e if jsonReq, ok := msg.(*jsonrpc.Request); ok && jsonReq.IsCall() { forCall = jsonReq } - // TODO: should we cancel this logical SSE request if/when jsonReq is canceled? - go c.handleSSE(requestSummary, resp, false, forCall) + // Handle the resulting stream. Note that ctx comes from the call, and + // therefore is already cancelled when the JSON-RPC request is cancelled + // (or rather, context cancellation is what *triggers* JSON-RPC + // cancellation) + go c.handleSSE(ctx, requestSummary, resp, forCall) default: resp.Body.Close() @@ -1500,23 +2425,65 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e return nil } -// testAuth controls whether a fake Authorization header is added to outgoing requests. -// TODO: replace with a better mechanism when client-side auth is in place. -var testAuth atomic.Bool - -func (c *streamableClientConn) setMCPHeaders(req *http.Request) { +func (c *streamableClientConn) setMCPHeaders(req *http.Request, msg jsonrpc.Message) error { c.mu.Lock() defer c.mu.Unlock() - if c.initializedResult != nil { + if c.oauthHandler != nil { + ts, err := c.oauthHandler.TokenSource(c.ctx) + if err != nil { + return err + } + if ts != nil { + token, err := ts.Token() + if err != nil { + // If the error is an invalid_grant oauth2.RetrieveError it indicates + // that the token source doesn't have valid authorization for the token + // endpoint, per RFC 6749 section 5.2. For example, the refresh token + // may be expired or invalid. + // + // In that case, ignore the error, skip setting the Authorization + // header, and proceed with the request. Callers that support + // authorization flows get a 401/403 response and trigger the + // Authorize() flow to refresh their token. + var retrieveErr *oauth2.RetrieveError + if !errors.As(err, &retrieveErr) || retrieveErr.ErrorCode != "invalid_grant" { + return err + } + } else if token != nil { + req.Header.Set("Authorization", "Bearer "+token.AccessToken) + } + } + } + if pv := protocolVersionFromMessage(msg); pv != "" { + req.Header.Set(protocolVersionHeader, pv) + } else if c.initializedResult != nil { req.Header.Set(protocolVersionHeader, c.initializedResult.ProtocolVersion) + } else if pv := protocolVersionFromContext(req.Context()); pv != "" { + req.Header.Set(protocolVersionHeader, pv) } if c.sessionID != "" { req.Header.Set(sessionIDHeader, c.sessionID) } - if testAuth.Load() { - req.Header.Set("Authorization", "Bearer foo") + + return nil +} + +// protocolVersionFromMessage recovers the SEP-2575 `_meta.protocolVersion` +// value from an outgoing JSON-RPC request, if present. It returns "" for +// notifications, responses, requests without a `_meta.protocolVersion`, or a +// nil msg. +func protocolVersionFromMessage(msg jsonrpc.Message) string { + req, ok := msg.(*jsonrpc.Request) + if !ok || req == nil { + return "" } + meta := extractRequestMeta(req.Params) + if meta == nil { + return "" + } + v, _ := meta[MetaKeyProtocolVersion].(string) + return v } func (c *streamableClientConn) handleJSON(requestSummary string, resp *http.Response) { @@ -1542,84 +2509,130 @@ func (c *streamableClientConn) handleJSON(requestSummary string, resp *http.Resp // persistent (for the main GET listener) or temporary (for a POST response). // // If forCall is set, it is the call that initiated the stream, and the -// stream is complete when we receive its response. -func (c *streamableClientConn) handleSSE(requestSummary string, initialResp *http.Response, persistent bool, forCall *jsonrpc2.Request) { - resp := initialResp - var lastEventID string - for { - // TODO: we should set a reasonable limit on the number of times we'll try - // getting a response for a given request. - // - // Eventually, if we don't get the response, we should stop trying and - // fail the request. - if resp != nil { - eventID, clientClosed := c.processStream(requestSummary, resp, forCall) - lastEventID = eventID +// stream is complete when we receive its response. Otherwise, this is the +// standalone stream. +func (c *streamableClientConn) handleSSE(ctx context.Context, requestSummary string, resp *http.Response, forCall *jsonrpc2.Request) { + // Track the last event ID to detect progress. + // The retry counter is only reset when progress is made (lastEventID advances). + // This prevents infinite retry loops when a server repeatedly terminates + // connections without making progress (#679). + var prevLastEventID string + retriesWithoutProgress := 0 - // If the connection was closed by the client, we're done. - if clientClosed { - return - } - // If the stream has ended, then do not reconnect if the stream is - // temporary (POST initiated SSE). - if lastEventID == "" && !persistent { + for { + lastEventID, reconnectDelay, clientClosed := c.processStream(ctx, requestSummary, resp, forCall) + + // If the connection was closed by the client, we're done. + if clientClosed { + return + } + // If we don't have a last event ID, we can never get the call response, so + // there's nothing to resume. For the standalone stream, we can reconnect, + // but we may just miss messages. + if lastEventID == "" && forCall != nil { + return + } + + // Check if we made progress (lastEventID advanced). + // Only reset the retry counter when actual progress is made. + if lastEventID != "" && lastEventID != prevLastEventID { + // Progress was made: reset the retry counter. + retriesWithoutProgress = 0 + prevLastEventID = lastEventID + } else { + // No progress: increment the retry counter. + retriesWithoutProgress++ + if retriesWithoutProgress > c.maxRetries { + if ctx.Err() == nil { + c.fail(fmt.Errorf("%s: exceeded %d retries without progress (session ID: %v)", requestSummary, c.maxRetries, c.sessionID)) + } return } } // The stream was interrupted or ended by the server. Attempt to reconnect. - newResp, err := c.reconnect(lastEventID) + newResp, err := c.connectSSE(ctx, lastEventID, reconnectDelay, false) if err != nil { - // All reconnection attempts failed: fail the connection. - c.fail(fmt.Errorf("%s: failed to reconnect (session ID: %v): %v", requestSummary, c.sessionID, err)) - return - } - resp = newResp - if resp.StatusCode == http.StatusMethodNotAllowed && persistent { - // [§2.2.3]: "The server MUST either return Content-Type: - // text/event-stream in response to this HTTP GET, or else return HTTP - // 405 Method Not Allowed, indicating that the server does not offer an - // SSE stream at this endpoint." - // - // [§2.2.3]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#listening-for-messages-from-the-server - - // The server doesn't support the standalone SSE stream. - resp.Body.Close() - return - } - if resp.StatusCode == http.StatusNotFound && persistent && !c.strict { - // modelcontextprotocol/gosdk#393: some servers return NotFound instead - // of MethodNotAllowed for the standalone SSE stream. - // - // Treat this like MethodNotAllowed in non-strict mode. - if c.logger != nil { - c.logger.Warn("got 404 instead of 405 for standalonw SSE stream") + // If the client didn't cancel this request, any failure to execute it + // breaks the logical MCP session. + if ctx.Err() == nil { + // All reconnection attempts failed: fail the connection. + c.fail(fmt.Errorf("%s: failed to reconnect (session ID: %v): %v", requestSummary, c.sessionID, err)) } - resp.Body.Close() return } - // (see equivalent handling in [streamableClientConn.Write]). - if resp.StatusCode == http.StatusNotFound { - c.fail(fmt.Errorf("%s: failed to reconnect (session ID: %v): %w", requestSummary, c.sessionID, errSessionMissing)) + + resp = newResp + if err := c.checkResponse(ctx, requestSummary, resp); err != nil { + c.fail(err) return } - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - resp.Body.Close() - c.fail(fmt.Errorf("%s: failed to reconnect: %v", requestSummary, http.StatusText(resp.StatusCode))) - return - } - // Reconnection was successful. Continue the loop with the new response. } } +// checkResponse checks the status code of the provided response, and +// translates it into an error if the request was unsuccessful. +// +// The response body is close if a non-nil error is returned. +func (c *streamableClientConn) checkResponse(ctx context.Context, requestSummary string, resp *http.Response) (err error) { + defer func() { + if err != nil { + resp.Body.Close() + } + }() + // Transient server errors (502, 503, 504, 429) should not break the connection. + // Wrap them with ErrRejected so the jsonrpc2 layer doesn't set writeErr. + if isTransientHTTPStatus(resp.StatusCode) { + return fmt.Errorf("%w: %s: %v", jsonrpc2.ErrRejected, requestSummary, http.StatusText(resp.StatusCode)) + } + // By default, always try to decode the body and surface the underlying + // JSON-RPC error, wrapping it with ErrRejected to prevent the connection from closing. + // Setting MCPGODEBUG=noprotocolerrorbody=1 restores the previous behavior. + if noprotocolerrorbody != "1" && (resp.StatusCode < 200 || resp.StatusCode >= 300) { + body, _ := io.ReadAll(resp.Body) + msg, _ := jsonrpc.DecodeMessage(body) + if response, ok := msg.(*jsonrpc.Response); ok && response.Error != nil { + return fmt.Errorf("%s: %w: %w: %v", requestSummary, response.Error, jsonrpc2.ErrRejected, http.StatusText(resp.StatusCode)) + } + } + // §2.5.3: "The server MAY terminate the session at any time, after + // which it MUST respond to requests containing that session ID with HTTP + // 404 Not Found." + // Only a stateful connection with a known session ID can experience ErrSessionMissing. + if sessionID := c.SessionID(); resp.StatusCode == http.StatusNotFound && sessionID != "" { + // Return an ErrSessionMissing to avoid sending a redundant DELETE when the + // session is already gone. + return fmt.Errorf("%s: failed to connect (session ID: %v): %w", requestSummary, sessionID, ErrSessionMissing) + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return fmt.Errorf("%s: %v", requestSummary, http.StatusText(resp.StatusCode)) + } + return nil +} + // processStream reads from a single response body, sending events to the // incoming channel. It returns the ID of the last processed event and a flag // indicating if the connection was closed by the client. If resp is nil, it // returns "", false. -func (c *streamableClientConn) processStream(requestSummary string, resp *http.Response, forCall *jsonrpc.Request) (lastEventID string, clientClosed bool) { - defer resp.Body.Close() - for evt, err := range scanEvents(resp.Body) { +func (c *streamableClientConn) processStream(ctx context.Context, requestSummary string, resp *http.Response, forCall *jsonrpc.Request) (lastEventID string, reconnectDelay time.Duration, clientClosed bool) { + defer func() { + // Drain any remaining unprocessed body. This allows the connection to be re-used after closing. + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + }() + for evt, err := range scanEventsLimited(resp.Body, c.maxEventSize) { if err != nil { + if ctx.Err() != nil { + return "", 0, true // don't reconnect: client cancelled + } + + // Malformed events are hard errors that indicate corrupted data or protocol + // violations. These should fail the connection permanently. + if errors.Is(err, errMalformedEvent) { + c.fail(fmt.Errorf("%s: %v", requestSummary, err)) + return "", 0, true + } + break } @@ -1627,30 +2640,55 @@ func (c *streamableClientConn) processStream(requestSummary string, resp *http.R lastEventID = evt.ID } + if evt.Retry != "" { + if n, err := strconv.ParseInt(evt.Retry, 10, 64); err == nil { + reconnectDelay = time.Duration(n) * time.Millisecond + } + } + + // According to SSE specification + // (https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation) + // events with an empty data buffer are allowed. + // In MCP these can be priming events (SEP-1699) that carry only a Last-Event-ID for stream resumption. + if len(evt.Data) == 0 { + continue + } + + // According to SSE spec, events with no name default to "message" + if evt.Name != "" && evt.Name != "message" { + continue + } + msg, err := jsonrpc.DecodeMessage(evt.Data) if err != nil { c.fail(fmt.Errorf("%s: failed to decode event: %v", requestSummary, err)) - return "", true + return "", 0, true } select { case c.incoming <- msg: + // Check if this is the response to our call, which terminates the request. + // (it could also be a server->client request or notification). if jsonResp, ok := msg.(*jsonrpc.Response); ok && forCall != nil { // TODO: we should never get a response when forReq is nil (the standalone SSE request). // We should detect this case. if jsonResp.ID == forCall.ID { - return "", true + return "", 0, true } } + case <-c.done: // The connection was closed by the client; exit gracefully. - return "", true + return "", 0, true } } // The loop finished without an error, indicating the server closed the stream. // // If the lastEventID is "", the stream is not retryable and we should // report a synthetic error for the call. + // + // Note that this is different from the cancellation case above, since the + // caller is still waiting for a response that will never come. if lastEventID == "" && forCall != nil { errmsg := &jsonrpc2.Response{ ID: forCall.ID, @@ -1661,57 +2699,98 @@ func (c *streamableClientConn) processStream(requestSummary string, resp *http.R case <-c.done: } } - return lastEventID, false + return lastEventID, reconnectDelay, false } -// reconnect handles the logic of retrying a connection with an exponential -// backoff strategy. It returns a new, valid HTTP response if successful, or -// an error if all retries are exhausted. -func (c *streamableClientConn) reconnect(lastEventID string) (*http.Response, error) { +// connectSSE handles the logic of connecting a text/event-stream connection. +// +// If lastEventID is set, it is the last-event ID of a stream being resumed. +// +// If connection fails, connectSSE retries with an exponential backoff +// strategy. It returns a new, valid HTTP response if successful, or an error +// if all retries are exhausted. +// +// reconnectDelay is the delay set by the server using the SSE retry field, or +// 0. +// +// If initial is set, this is the initial attempt. +// +// If connectSSE exits due to context cancellation, the result is (nil, ctx.Err()). +func (c *streamableClientConn) connectSSE(ctx context.Context, lastEventID string, reconnectDelay time.Duration, initial bool) (*http.Response, error) { var finalErr error - - // We can reach the 'reconnect' path through the standlone SSE request, in which case - // lastEventID will be "". - // - // In this case, we need an initial attempt. attempt := 0 - if lastEventID != "" { + if !initial { + // We've already connected successfully once, so delay subsequent + // reconnections. Otherwise, if the server returns 200 but terminates the + // connection, we'll reconnect as fast as we can, ad infinitum. + // + // TODO: we should consider also setting a limit on total attempts for one + // logical request. attempt = 1 } - + delay := calculateReconnectDelay(attempt) + if reconnectDelay > 0 { + delay = reconnectDelay // honor the server's requested initial delay + } for ; attempt <= c.maxRetries; attempt++ { select { case <-c.done: return nil, fmt.Errorf("connection closed by client during reconnect") - case <-time.After(calculateReconnectDelay(attempt)): - resp, err := c.establishSSE(lastEventID) + + case <-ctx.Done(): + // If the connection context is canceled, the request below will not + // succeed anyway. + return nil, ctx.Err() + + case <-time.After(delay): + req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.url, nil) + if err != nil { + return nil, err + } + if err := c.setMCPHeaders(req, nil); err != nil { + return nil, err + } + if lastEventID != "" { + req.Header.Set(lastEventIDHeader, lastEventID) + } + req.Header.Set("Accept", "text/event-stream") + resp, err := c.client.Do(req) if err != nil { finalErr = err // Store the error and try again. + delay = calculateReconnectDelay(attempt + 1) continue } return resp, nil } } - // If the loop completes, all retries have failed. + // If the loop completes, all retries have failed, or the client is closing. if finalErr != nil { return nil, fmt.Errorf("connection failed after %d attempts: %w", c.maxRetries, finalErr) } - return nil, fmt.Errorf("connection failed after %d attempts", c.maxRetries) + return nil, fmt.Errorf("connection aborted after %d attempts", c.maxRetries) } // Close implements the [Connection] interface. func (c *streamableClientConn) Close() error { c.closeOnce.Do(func() { - if errors.Is(c.failure(), errSessionMissing) { + if errors.Is(c.failure(), ErrSessionMissing) { // If the session is missing, no need to delete it. + } else if c.SessionID() == "" { + // No session was established (e.g. the server is stateless), + // so there is nothing to delete. } else { - req, err := http.NewRequestWithContext(c.ctx, http.MethodDelete, c.url, nil) + reqCtx, stop := context.WithTimeout(c.ctx, closeDeleteTimeout) + defer stop() + req, err := http.NewRequestWithContext(reqCtx, http.MethodDelete, c.url, nil) if err != nil { c.closeErr = err } else { - c.setMCPHeaders(req) - if _, err := c.client.Do(req); err != nil { + if err := c.setMCPHeaders(req, nil); err != nil { c.closeErr = err + } else if resp, err := c.client.Do(req); err != nil { + c.closeErr = err + } else { + resp.Body.Close() } } } @@ -1723,30 +2802,13 @@ func (c *streamableClientConn) Close() error { return c.closeErr } -// establishSSE establishes the persistent SSE listening stream. -// It is used for reconnect attempts using the Last-Event-ID header to -// resume a broken stream where it left off. -func (c *streamableClientConn) establishSSE(lastEventID string) (*http.Response, error) { - req, err := http.NewRequestWithContext(c.ctx, http.MethodGet, c.url, nil) - if err != nil { - return nil, err - } - c.setMCPHeaders(req) - if lastEventID != "" { - req.Header.Set("Last-Event-ID", lastEventID) - } - req.Header.Set("Accept", "text/event-stream") - - return c.client.Do(req) -} - // calculateReconnectDelay calculates a delay using exponential backoff with full jitter. func calculateReconnectDelay(attempt int) time.Duration { if attempt == 0 { return 0 } // Calculate the exponential backoff using the grow factor. - backoffDuration := time.Duration(float64(reconnectInitialDelay) * math.Pow(reconnectGrowFactor, float64(attempt-1))) + backoffDuration := time.Duration(float64(reconnectInitialDelay.Load()) * math.Pow(reconnectGrowFactor, float64(attempt-1))) // Cap the backoffDuration at maxDelay. backoffDuration = min(backoffDuration, reconnectMaxDelay) @@ -1755,3 +2817,17 @@ func calculateReconnectDelay(attempt int) time.Duration { return backoffDuration + jitter } + +// isTransientHTTPStatus reports whether the HTTP status code indicates a +// transient server error that should not permanently break the connection. +func isTransientHTTPStatus(statusCode int) bool { + switch statusCode { + case http.StatusInternalServerError, // 500 + http.StatusBadGateway, // 502 + http.StatusServiceUnavailable, // 503 + http.StatusGatewayTimeout, // 504 + http.StatusTooManyRequests: // 429 + return true + } + return false +} diff --git a/mcp/streamable_bench_test.go b/mcp/streamable_bench_test.go index bf82ee5..fc7b0ef 100644 --- a/mcp/streamable_bench_test.go +++ b/mcp/streamable_bench_test.go @@ -5,10 +5,15 @@ package mcp_test import ( - "context" + "flag" + "log" "net/http" "net/http/httptest" + "os" "reflect" + "runtime" + "runtime/pprof" + "strings" "testing" "github.com/google/jsonschema-go/jsonschema" @@ -45,8 +50,7 @@ func BenchmarkStreamableServing(b *testing.B) { httpServer := httptest.NewServer(handler) defer httpServer.Close() - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := b.Context() session, err := mcp.NewClient(testImpl, nil).Connect(ctx, &mcp.StreamableClientTransport{Endpoint: httpServer.URL}, nil) if err != nil { b.Fatal(err) @@ -65,3 +69,61 @@ func BenchmarkStreamableServing(b *testing.B) { } } } + +var streamableHeap = flag.String("streamable_heap", "", "if set, write streamable heap profiles with this prefix") + +func BenchmarkStreamableServing_BadSessions(b *testing.B) { + server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + + handler := mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { + return server + }, &mcp.StreamableHTTPOptions{JSONResponse: true}) + + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + ctx := b.Context() + + if *streamableHeap != "" { + writeHeap := func(file string) { + // GC a couple times to ensure accurate heap. + runtime.GC() + runtime.GC() + f, err := os.Create(file) + if err != nil { + log.Fatal("could not create memory profile: ", err) + } + defer func() { + if err := f.Close(); err != nil { + b.Errorf("writing heap file %q: %v", file, err) + } + }() + if err := pprof.Lookup("heap").WriteTo(f, 0); err != nil { + b.Errorf("could not write heap profile: %v", err) + } + } + writeHeap(*streamableHeap + ".before") + defer writeHeap(*streamableHeap + ".after") + } + + b.ResetTimer() + for range b.N { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, httpServer.URL, strings.NewReader("{}")) + if err != nil { + b.Fatal(err) + } + req.Header.Add("Accept", "application/json") + req.Header.Add("Accept", "text/event-stream") + resp, err := http.DefaultClient.Do(req) + if err != nil { + b.Fatal(err) + } + if got, want := resp.StatusCode, http.StatusBadRequest; got != want { + b.Fatalf("POST got status %d, want %d", got, want) + } + if got := resp.Header.Get("Mcp-Session-Id"); got != "" { + b.Fatalf("POST got unexpected session ID") + } + resp.Body.Close() + } +} diff --git a/mcp/streamable_client.go b/mcp/streamable_client.go new file mode 100644 index 0000000..74940b3 --- /dev/null +++ b/mcp/streamable_client.go @@ -0,0 +1,231 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +// TODO: move client-side streamable HTTP logic from streamable.go to this file. + +package mcp + +/* +Streamable HTTP Client Design + +This document describes the client-side implementation of the MCP streamable +HTTP transport, as defined by the MCP spec: +https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#streamable-http + +# Overview + +The client-side streamable transport allows an MCP client to communicate with a +server over HTTP, sending messages via POST and receiving responses via either +JSON or server-sent events (SSE). The implementation consists of two main +components: + + ┌─────────────────────────────────────────────────────────────────┐ + │ [StreamableClientTransport] │ + │ Transport configuration; creates connections via Connect() │ + └─────────────────────────────────────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────────────────────────────────────┐ + │ [streamableClientConn] │ + │ Connection implementation; handles HTTP request/response │ + └─────────────────────────────────────────────────────────────────┘ + │ + ├──────────────────────────────────────┐ + ▼ ▼ + ┌─────────────────────────────────────────┐ ┌────────────────────────────────────┐ + │ POST request handlers │ │ Standalone SSE stream │ + │ (one per outgoing message/call) │ │ (server-initiated messages) │ + └─────────────────────────────────────────┘ └────────────────────────────────────┘ + +# Sessions + +The client optionally maintains a session with the server, identified by a +session ID (Mcp-Session-Id header). Sessionless servers (those that never +send Mcp-Session-Id) are fully supported; the client simply omits the +header from subsequent requests and skips the DELETE on close. + +When a session is present: + + - Session ID is received from the server after initialization + - Client includes the session ID in all subsequent requests + - Session ends when the client calls Close() (sends DELETE) or server returns 404 + +[streamableClientConn] stores the session state: + - [streamableClientConn.sessionID]: Server-assigned session identifier + - [streamableClientConn.initializedResult]: Protocol version and server capabilities + +# Connection Lifecycle + +1. Connect: [StreamableClientTransport.Connect] creates a [streamableClientConn] + with a detached context for the connection's lifetime. The context is detached + to prevent the standalone SSE stream from being cancelled when the original + Connect context times out. + +2. Initialize: The MCP client sends initialize/initialized messages. Upon + receiving [InitializeResult], the connection: + - Stores the negotiated protocol version for the Mcp-Protocol-Version header + - Captures the session ID from the Mcp-Session-Id response header + - Starts the standalone SSE stream via [streamableClientConn.connectStandaloneSSE] + +3. Operation: Messages are sent via POST, responses received via JSON or SSE. + +4. Close: [streamableClientConn.Close] sends a DELETE request to terminate + the session (unless the session is already gone), then cancels the connection + context to clean up the standalone SSE stream. + +# Sending Messages (Write) + +[streamableClientConn.Write] sends all outgoing messages via HTTP POST: + + POST /endpoint + Content-Type: application/json + Accept: application/json, text/event-stream + Mcp-Protocol-Version: + Mcp-Session-Id: + + + +The server may respond with: + - 202 Accepted: Message received, no response body (notifications/responses) + - 200 OK with application/json: Single JSON-RPC response + - 200 OK with text/event-stream: SSE stream of responses + +# Receiving Messages (Read) + +[streamableClientConn.Read] returns messages from the [streamableClientConn.incoming] +channel, which is populated by multiple concurrent goroutines: + +1. POST response handlers ([streamableClientConn.handleJSON] and + [streamableClientConn.handleSSE]): Process responses from POST requests + +2. Standalone SSE stream: Receives server-initiated requests and notifications + +The client handles both response formats: + - JSON: [streamableClientConn.handleJSON] reads body, decodes message + - SSE: [streamableClientConn.handleSSE] scans events, decodes each message + +# Standalone SSE Stream + +After initialization, [streamableClientConn.sessionUpdated] triggers +[streamableClientConn.connectStandaloneSSE] to open a GET request for +server-initiated messages: + + GET /endpoint + Accept: text/event-stream + Mcp-Session-Id: + +Stream behavior: + - Optional: Server may return 405 Method Not Allowed (spec-compliant) or + other 4xx errors (tolerated in non-strict mode for compatibility) + - Persistent: Runs for the connection lifetime in a background goroutine + - Resumable: Uses Last-Event-ID header on reconnection if server provides event IDs + - Reconnects: Automatic reconnection with exponential backoff on interruption + +# Stream Resumption + +When an SSE stream (standalone or POST response) is interrupted, the client +attempts to reconnect using [streamableClientConn.connectSSE]: + +Event ID tracking: + - [streamableClientConn.processStream] tracks the last received event ID + - On reconnection, the Last-Event-ID header is set to resume from that point + - Server replays missed events if it has an [EventStore] configured + +See [calculateReconnectDelay] for the reconnect delay details. + +Server-initiated reconnection (SEP-1699) + - SSE retry field: Sets the delay for the next reconnect attempt + - If server doesn't provide event IDs, non-standalone streams don't reconnect + +# Response Formats + +The client must handle two response formats from POST requests: + +1. application/json: Single JSON-RPC response + - Body contains one JSON-RPC message + - Handled by [streamableClientConn.handleJSON] + - Simpler but doesn't support streaming or server-initiated messages + +2. text/event-stream: SSE stream of messages + - Body contains SSE events with JSON-RPC messages + - Handled by [streamableClientConn.handleSSE] + - Supports multiple messages and server-initiated communication + - Stream completes when the response to the originating call is received + +# HTTP Methods + + - POST: Send JSON-RPC messages (requests, responses, notifications) + - Used by [streamableClientConn.Write] + - Response may be JSON or SSE + + - GET: Open or resume SSE stream for server-initiated messages + - Used by [streamableClientConn.connectSSE] + - Always expects text/event-stream response (or 405) + + - DELETE: Terminate the session + - Used by [streamableClientConn.Close] + - Skipped if session is already known to be gone ([ErrSessionMissing]) + or if no session was established (sessionless server) + +# Error Handling + +Errors are categorized and handled differently: + +1. Transient (recoverable via reconnection): + - Network interruption during SSE streaming + - Connection reset or timeout + - Triggers reconnection in [streamableClientConn.handleSSE] + +2. Terminal (breaks the connection): + - 404 Not Found: Session terminated by server ([ErrSessionMissing]) + - Message decode errors: Protocol violation + - Context cancellation: Client closed connection + - Mismatched session IDs: Protocol error (only relevant for servers that use sessions) + - See issue #683: our terminal errors are too strict. + +Terminal errors are stored via [streamableClientConn.fail] and returned by +subsequent [streamableClientConn.Read] calls. The [streamableClientConn.failed] +channel signals that the connection is broken. + +Special case: [ErrSessionMissing] indicates the server has terminated the session, +so [streamableClientConn.Close] skips the DELETE request. + +# Protocol Version Header + +After initialization, all requests include: + + Mcp-Protocol-Version: + +This header (set by [streamableClientConn.setMCPHeaders]): + - Allows the server to handle requests per the negotiated protocol + - Is omitted before initialization completes + - Uses the version from [streamableClientConn.initializedResult] + +# Key Implementation Details + +[StreamableClientTransport] configuration: + - [StreamableClientTransport.Endpoint]: URL of the MCP server + - [StreamableClientTransport.HTTPClient]: Custom HTTP client (optional) + - [StreamableClientTransport.MaxRetries]: Reconnection attempts (default 5) + +[streamableClientConn] handles the [Connection] interface: + - [streamableClientConn.Read]: Returns messages from incoming channel + - [streamableClientConn.Write]: Sends messages via POST, starts response handlers + - [streamableClientConn.Close]: Sends DELETE, cancels context, closes done channel + +State management: + - [streamableClientConn.incoming]: Buffered channel for received messages + - [streamableClientConn.sessionID]: Server-assigned session identifier (empty for sessionless servers) + - [streamableClientConn.initializedResult]: Cached for protocol version header + - [streamableClientConn.failed]: Channel closed on terminal error + - [streamableClientConn.done]: Channel closed on graceful shutdown + - [streamableClientConn.ctx]: Detached context for connection lifetime + - [streamableClientConn.cancel]: Cancels ctx to terminate SSE streams + +Context handling: + - Connection context is detached from [StreamableClientTransport.Connect] context + using [xcontext.Detach] to preserve context values (for auth middleware) while + preventing premature cancellation of the standalone SSE stream + - Individual POST requests use caller-provided contexts for cancellation +*/ diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index 4a4f5c6..77ff81a 100644 --- a/mcp/streamable_client_test.go +++ b/mcp/streamable_client_test.go @@ -6,34 +6,43 @@ package mcp import ( "context" + "encoding/json" + "errors" "fmt" "io" "net/http" "net/http/httptest" "strings" "sync" + "sync/atomic" "testing" "time" "github.com/google/go-cmp/cmp" + "github.com/modelcontextprotocol/go-sdk/auth" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" "github.com/modelcontextprotocol/go-sdk/jsonrpc" + "golang.org/x/oauth2" ) type streamableRequestKey struct { httpMethod string // http method sessionID string // session ID header jsonrpcMethod string // jsonrpc method, or "" for non-requests + lastEventID string // Last-Event-ID header } type header map[string]string +// TODO: replace body and status fields with responseFunc; add helpers to reduce duplication. type streamableResponse struct { - header header // response headers - status int // or http.StatusOK - body string // or "" - optional bool // if set, request need not be sent - wantProtocolVersion string // if "", unchecked + header header // response headers + status int // or http.StatusOK; ignored if responseFunc is set + body string // or ""; ignored if responseFunc is set + responseFunc func(r *jsonrpc.Request) (string, int) // if set, overrides body and status + optional bool // if set, request need not be sent + wantProtocolVersion string // if "", unchecked + done chan struct{} // if set, receive from this channel before terminating the request } type fakeResponses map[streamableRequestKey]*streamableResponse @@ -42,17 +51,17 @@ type fakeStreamableServer struct { t *testing.T responses fakeResponses - callMu sync.Mutex - calls map[streamableRequestKey]int + calledMu sync.Mutex + called map[streamableRequestKey]bool } func (s *fakeStreamableServer) missingRequests() []streamableRequestKey { - s.callMu.Lock() - defer s.callMu.Unlock() + s.calledMu.Lock() + defer s.calledMu.Unlock() var unused []streamableRequestKey for k, resp := range s.responses { - if s.calls[k] == 0 && !resp.optional { + if !s.called[k] && !resp.optional { unused = append(unused, k) } } @@ -61,9 +70,11 @@ func (s *fakeStreamableServer) missingRequests() []streamableRequestKey { func (s *fakeStreamableServer) ServeHTTP(w http.ResponseWriter, req *http.Request) { key := streamableRequestKey{ - httpMethod: req.Method, - sessionID: req.Header.Get(sessionIDHeader), + httpMethod: req.Method, + sessionID: req.Header.Get(sessionIDHeader), + lastEventID: req.Header.Get(lastEventIDHeader), } + var jsonrpcReq *jsonrpc.Request if req.Method == http.MethodPost { body, err := io.ReadAll(req.Body) if err != nil { @@ -79,35 +90,67 @@ func (s *fakeStreamableServer) ServeHTTP(w http.ResponseWriter, req *http.Reques } if r, ok := msg.(*jsonrpc.Request); ok { key.jsonrpcMethod = r.Method + jsonrpcReq = r } } - s.callMu.Lock() - if s.calls == nil { - s.calls = make(map[streamableRequestKey]int) + s.calledMu.Lock() + if s.called == nil { + s.called = make(map[streamableRequestKey]bool) } - s.calls[key]++ - s.callMu.Unlock() + s.called[key] = true + s.calledMu.Unlock() resp, ok := s.responses[key] if !ok { + if key.jsonrpcMethod == "server/discover" { + // Return MethodNotFound to trigger fallback to legacy initialize. + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + respErr := &jsonrpc.Error{ + Code: jsonrpc.CodeMethodNotFound, + Message: `method not found: "server/discover"`, + } + var id jsonrpc.ID + if jsonrpcReq != nil { + id = jsonrpcReq.ID + } + respMsg, _ := jsonrpc2.NewResponse(id, nil, respErr) + data, _ := jsonrpc2.EncodeMessage(respMsg) + w.Write(data) + return + } s.t.Errorf("missing response for %v", key) http.Error(w, "no response", http.StatusInternalServerError) return } - for k, v := range resp.header { - w.Header().Set(k, v) - } + + // Determine body and status, potentially using responseFunc for dynamic responses. + body := resp.body status := resp.status + if resp.responseFunc != nil { + body, status = resp.responseFunc(jsonrpcReq) + } if status == 0 { status = http.StatusOK } + + for k, v := range resp.header { + w.Header().Set(k, v) + } + rc := http.NewResponseController(w) w.WriteHeader(status) + rc.Flush() // flush response headers if v := req.Header.Get(protocolVersionHeader); v != resp.wantProtocolVersion && resp.wantProtocolVersion != "" { s.t.Errorf("%v: bad protocol version header: got %q, want %q", key, v, resp.wantProtocolVersion) } - w.Write([]byte(resp.body)) + w.Write([]byte(body)) + rc.Flush() // flush response + + if resp.done != nil { + <-resp.done + } } var ( @@ -117,7 +160,8 @@ var ( Logging: &LoggingCapabilities{}, Tools: &ToolCapabilities{ListChanged: true}, }, - ProtocolVersion: latestProtocolVersion, + // Pin negotiated version to 2025-11-25 + ProtocolVersion: protocolVersion20251125, ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"}, } initResp = resp(1, initResult, nil) @@ -141,25 +185,24 @@ func TestStreamableClientTransportLifecycle(t *testing.T) { fake := &fakeStreamableServer{ t: t, responses: fakeResponses{ - {"POST", "", methodInitialize}: { + {"POST", "", methodInitialize, ""}: { header: header{ "Content-Type": "application/json", sessionIDHeader: "123", }, body: jsonBody(t, initResp), }, - {"POST", "123", notificationInitialized}: { + {"POST", "123", notificationInitialized, ""}: { status: http.StatusAccepted, - wantProtocolVersion: latestProtocolVersion, + wantProtocolVersion: protocolVersion20251125, }, - {"GET", "123", ""}: { + {"GET", "123", "", ""}: { header: header{ "Content-Type": "text/event-stream", }, - optional: true, - wantProtocolVersion: latestProtocolVersion, + wantProtocolVersion: protocolVersion20251125, }, - {"DELETE", "123", ""}: {}, + {"DELETE", "123", "", ""}: {}, }, } @@ -168,7 +211,10 @@ func TestStreamableClientTransportLifecycle(t *testing.T) { transport := &StreamableClientTransport{Endpoint: httpServer.URL} client := NewClient(testImpl, nil) - session, err := client.Connect(ctx, transport, nil) + // Pin to 2025-11-25: the fixture's canned initialize response uses + // hardcoded id=1, which only matches when initialize is the first + // request. Under 2026-07-28 the client probes server/discover first. + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatalf("client.Connect() failed: %v", err) } @@ -193,22 +239,21 @@ func TestStreamableClientRedundantDelete(t *testing.T) { fake := &fakeStreamableServer{ t: t, responses: fakeResponses{ - {"POST", "", methodInitialize}: { + {"POST", "", methodInitialize, ""}: { header: header{ "Content-Type": "application/json", sessionIDHeader: "123", }, body: jsonBody(t, initResp), }, - {"POST", "123", notificationInitialized}: { + {"POST", "123", notificationInitialized, ""}: { status: http.StatusAccepted, - wantProtocolVersion: latestProtocolVersion, + wantProtocolVersion: protocolVersion20251125, }, - {"GET", "123", ""}: { - status: http.StatusMethodNotAllowed, - optional: true, + {"GET", "123", "", ""}: { + status: http.StatusMethodNotAllowed, }, - {"POST", "123", methodListTools}: { + {"POST", "123", methodListTools, ""}: { status: http.StatusNotFound, }, }, @@ -219,7 +264,10 @@ func TestStreamableClientRedundantDelete(t *testing.T) { transport := &StreamableClientTransport{Endpoint: httpServer.URL} client := NewClient(testImpl, nil) - session, err := client.Connect(ctx, transport, nil) + // Pin to 2025-11-25: the fixture's canned initialize response uses + // hardcoded id=1, which only matches when initialize is the first + // request. Under 2026-07-28 the client probes server/discover first. + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatalf("client.Connect() failed: %v", err) } @@ -233,50 +281,91 @@ func TestStreamableClientRedundantDelete(t *testing.T) { } } +// TestStreamableClientCloseUnresponsiveDelete checks that Close stays bounded +// when the peer accepts the session-termination DELETE but never answers it. +func TestStreamableClientCloseUnresponsiveDelete(t *testing.T) { + base := NewStreamableHTTPHandler(func(*http.Request) *Server { + return NewServer(testImpl, nil) + }, nil) + httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodDelete { + <-r.Context().Done() + return + } + base.ServeHTTP(w, r) + })) + t.Cleanup(func() { + // The DELETE handler only returns once its request is cancelled, so + // drop the connections before waiting for outstanding requests. + httpServer.CloseClientConnections() + httpServer.Close() + }) + + client := NewClient(testImpl, nil) + session, err := client.Connect(context.Background(), + &StreamableClientTransport{Endpoint: httpServer.URL}, nil) + if err != nil { + t.Fatalf("client.Connect() failed: %v", err) + } + if session.ID() == "" { + t.Fatal("empty session ID: Close would not send DELETE") + } + + done := make(chan struct{}) + go func() { + session.Close() + close(done) + }() + select { + case <-done: + case <-time.After(closeDeleteTimeout + 5*time.Second): + t.Fatal("Close() blocked on an unanswered DELETE") + } +} + func TestStreamableClientGETHandling(t *testing.T) { ctx := context.Background() tests := []struct { status int wantErrorContaining string + contentType string }{ - {http.StatusOK, ""}, - {http.StatusMethodNotAllowed, ""}, - {http.StatusBadRequest, "standalone SSE"}, + {http.StatusOK, "", "text/event-stream"}, + {http.StatusOK, "", "text/event-stream; charset=utf-8"}, + {http.StatusMethodNotAllowed, "", "text/event-stream"}, + //// The client error status code is not treated as an error in non-strict + //// mode. + {http.StatusNotFound, "", "text/event-stream"}, + {http.StatusBadRequest, "", "text/event-stream"}, + {http.StatusInternalServerError, "standalone SSE", "text/event-stream"}, + {http.StatusOK, "", "text/html; charset=utf-8"}, } for _, test := range tests { - t.Run(fmt.Sprintf("status=%d", test.status), func(t *testing.T) { + t.Run(fmt.Sprintf("status=%d content_type=%q", test.status, test.contentType), func(t *testing.T) { fake := &fakeStreamableServer{ t: t, responses: fakeResponses{ - {"POST", "", methodInitialize}: { + {"POST", "", methodInitialize, ""}: { header: header{ "Content-Type": "application/json; charset=utf-8", // should ignore the charset sessionIDHeader: "123", }, body: jsonBody(t, initResp), }, - {"POST", "123", notificationInitialized}: { + {"POST", "123", notificationInitialized, ""}: { status: http.StatusAccepted, - wantProtocolVersion: latestProtocolVersion, + wantProtocolVersion: protocolVersion20251125, }, - {"GET", "123", ""}: { + {"GET", "123", "", ""}: { header: header{ - "Content-Type": "text/event-stream", + "Content-Type": test.contentType, }, status: test.status, - wantProtocolVersion: latestProtocolVersion, + wantProtocolVersion: protocolVersion20251125, }, - {"POST", "123", methodListTools}: { - header: header{ - "Content-Type": "application/json", - sessionIDHeader: "123", - }, - body: jsonBody(t, resp(2, &ListToolsResult{Tools: []*Tool{}}, nil)), - optional: true, - }, - {"DELETE", "123", ""}: {optional: true}, + {"DELETE", "123", "", ""}: {optional: true}, }, } httpServer := httptest.NewServer(fake) @@ -284,37 +373,21 @@ func TestStreamableClientGETHandling(t *testing.T) { transport := &StreamableClientTransport{Endpoint: httpServer.URL} client := NewClient(testImpl, nil) - session, err := client.Connect(ctx, transport, nil) - if err != nil { - t.Fatalf("client.Connect() failed: %v", err) + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ + ProtocolVersion: protocolVersion20251125, + }) + if err == nil { + defer session.Close() } - - // Since we need the client to observe the result of the hanging GET, - // wait for all requests to be handled. - start := time.Now() - delay := 1 * time.Millisecond - for range 10 { - if len(fake.missingRequests()) == 0 { - break + if test.wantErrorContaining != "" { + if err == nil { + t.Fatalf("Connect succeeded unexpectedly, want error containing %q", test.wantErrorContaining) + } + if got := err.Error(); !strings.Contains(got, test.wantErrorContaining) { + t.Errorf("Connect error = %q, want containing %q", got, test.wantErrorContaining) } - time.Sleep(delay) - delay *= 2 - } - if missing := fake.missingRequests(); len(missing) > 0 { - t.Errorf("did not receive expected requests after %s: %v", time.Since(start), missing) - } - - _, err = session.ListTools(ctx, nil) - if (err != nil) != (test.wantErrorContaining != "") { - t.Errorf("After initialization, got error %v, want containing %q", err, test.wantErrorContaining) } else if err != nil { - if !strings.Contains(err.Error(), test.wantErrorContaining) { - t.Errorf("After initialization, got error %s, want containing %q", err, test.wantErrorContaining) - } - } - - if err := session.Close(); err != nil { - t.Errorf("closing session: %v", err) + t.Fatalf("Connect failed: %v", err) } }) } @@ -329,38 +402,41 @@ func TestStreamableClientStrictness(t *testing.T) { initializedStatus int getStatus int wantConnectError bool - wantListError bool }{ - {"conformant server", true, http.StatusAccepted, http.StatusMethodNotAllowed, false, false}, - {"strict initialized", true, http.StatusOK, http.StatusMethodNotAllowed, true, false}, - {"unstrict initialized", false, http.StatusOK, http.StatusMethodNotAllowed, false, false}, - {"strict GET", true, http.StatusAccepted, http.StatusNotFound, false, true}, - {"unstrict GET", false, http.StatusOK, http.StatusNotFound, false, false}, + {"conformant server", true, http.StatusAccepted, http.StatusMethodNotAllowed, false}, + {"strict initialized", true, http.StatusOK, http.StatusMethodNotAllowed, true}, + {"unstrict initialized", false, http.StatusOK, http.StatusMethodNotAllowed, false}, + {"strict GET", true, http.StatusAccepted, http.StatusNotFound, true}, + // The client error status code is not treated as an error in non-strict + // mode. + {"unstrict GET on StatusNotFound", false, http.StatusOK, http.StatusNotFound, false}, + {"unstrict GET on StatusBadRequest", false, http.StatusOK, http.StatusBadRequest, false}, + {"GET on InternlServerError", false, http.StatusOK, http.StatusInternalServerError, true}, } for _, test := range tests { t.Run(test.label, func(t *testing.T) { fake := &fakeStreamableServer{ t: t, responses: fakeResponses{ - {"POST", "", methodInitialize}: { + {"POST", "", methodInitialize, ""}: { header: header{ "Content-Type": "application/json", sessionIDHeader: "123", }, body: jsonBody(t, initResp), }, - {"POST", "123", notificationInitialized}: { + {"POST", "123", notificationInitialized, ""}: { status: test.initializedStatus, - wantProtocolVersion: latestProtocolVersion, + wantProtocolVersion: protocolVersion20251125, }, - {"GET", "123", ""}: { + {"GET", "123", "", ""}: { header: header{ "Content-Type": "text/event-stream", }, status: test.getStatus, - wantProtocolVersion: latestProtocolVersion, + wantProtocolVersion: protocolVersion20251125, }, - {"POST", "123", methodListTools}: { + {"POST", "123", methodListTools, ""}: { header: header{ "Content-Type": "application/json", sessionIDHeader: "123", @@ -368,7 +444,7 @@ func TestStreamableClientStrictness(t *testing.T) { body: jsonBody(t, resp(2, &ListToolsResult{Tools: []*Tool{}}, nil)), optional: true, }, - {"DELETE", "123", ""}: {optional: true}, + {"DELETE", "123", "", ""}: {optional: true}, }, } httpServer := httptest.NewServer(fake) @@ -376,30 +452,20 @@ func TestStreamableClientStrictness(t *testing.T) { transport := &StreamableClientTransport{Endpoint: httpServer.URL, strict: test.strict} client := NewClient(testImpl, nil) - session, err := client.Connect(ctx, transport, nil) + // Pin to 2025-11-25: the fixture's canned initialize response + // uses hardcoded id=1, which only matches when initialize is + // the first request. Under 2026-07-28 the client probes + // server/discover first. + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if (err != nil) != test.wantConnectError { t.Errorf("client.Connect() returned error %v; want error: %t", err, test.wantConnectError) } if err != nil { return } - // Since we need the client to observe the result of the hanging GET, - // wait for all requests to be handled. - start := time.Now() - delay := 1 * time.Millisecond - for range 10 { - if len(fake.missingRequests()) == 0 { - break - } - time.Sleep(delay) - delay *= 2 - } - if missing := fake.missingRequests(); len(missing) > 0 { - t.Errorf("did not receive expected requests after %s: %v", time.Since(start), missing) - } _, err = session.ListTools(ctx, nil) - if (err != nil) != test.wantListError { - t.Errorf("ListTools returned error %v; want error: %t", err, test.wantListError) + if err != nil { + t.Errorf("ListTools failed: %v", err) } if err := session.Close(); err != nil { t.Errorf("closing session: %v", err) @@ -415,14 +481,14 @@ func TestStreamableClientUnresumableRequest(t *testing.T) { fake := &fakeStreamableServer{ t: t, responses: fakeResponses{ - {"POST", "", methodInitialize}: { + {"POST", "", methodInitialize, ""}: { header: header{ "Content-Type": "text/event-stream", sessionIDHeader: "123", }, body: "", }, - {"DELETE", "123", ""}: {optional: true}, + {"DELETE", "123", "", ""}: {optional: true}, }, } httpServer := httptest.NewServer(fake) @@ -430,7 +496,7 @@ func TestStreamableClientUnresumableRequest(t *testing.T) { transport := &StreamableClientTransport{Endpoint: httpServer.URL} client := NewClient(testImpl, nil) - cs, err := client.Connect(ctx, transport, nil) + cs, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err == nil { cs.Close() t.Fatalf("Connect succeeded unexpectedly") @@ -442,3 +508,1524 @@ func TestStreamableClientUnresumableRequest(t *testing.T) { t.Errorf("Connect: got error %v, want containing %q", err, msg) } } + +func TestStreamableClientResumption_Cancelled(t *testing.T) { + // This test verifies that the resumed requests are closed when their context + // is cancelled (issue #662). + + // This test (unfortunately) relies on timing, so may have false positives. + // + // Set the reconnect initial delay to some small(ish) value so that the test + // doesn't take too long. But this value must be large enough that we mostly + // avoid races in the tests below, where one test cases is intended to be in + // between the initial attempt and first reconnection. + // + // For easier tuning (and debugging), factor out the tick size. + // + // TODO(#680): experiment with instead using synctest. + const tick = 10 * time.Millisecond + defer func(delay int64) { + reconnectInitialDelay.Store(delay) + }(reconnectInitialDelay.Load()) + reconnectInitialDelay.Store(int64(2 * tick)) + + // The setup: terminate a request stream and make the resumed request hang + // indefinitely. CallTool should still exit when its context is canceled. + // + // This should work whether we're handling the initial request, waiting to + // retry, or handling the retry. + // + // Furthermore, closing the client connection should not hang, because there + // should be no ongoing requests. + + tests := []struct { + label string + cancelAfter time.Duration + }{ + {"in process", 1 * tick}, // cancel while the request is being handled + // initial request terminates at 2 ticks (see below) + {"awaiting retry", 3 * tick}, // cancel in-between first and second attempt + // retry starts at 4 ticks (=2+2) + {"in retry", 5 * tick}, // cancel while second attempt is hanging + } + + for _, test := range tests { + t.Run(test.label, func(t *testing.T) { + ctx := context.Background() + + // done will be closed when the test exits: used to simulate requests that + // hang indefinitely. + initialRequestDone := make(chan struct{}) // closed below + allDone := make(chan struct{}) + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + body: jsonBody(t, initResp), + }, + {"POST", "123", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + {"GET", "123", "", ""}: { + header: header{ + "Content-Type": "text/event-stream", + }, + status: http.StatusMethodNotAllowed, // don't allow the standalone stream + }, + {"POST", "123", methodCallTool, ""}: { + header: header{ + "Content-Type": "text/event-stream", + }, + status: http.StatusOK, + body: `id: 1 +data: { "jsonrpc": "2.0", "method": "notifications/message", "params": { "level": "error", "data": "bad" } } + +`, + done: initialRequestDone, + }, + {"POST", "123", methodListTools, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + body: jsonBody(t, resp(3, &ListToolsResult{Tools: []*Tool{}}, nil)), + }, + {"GET", "123", "", "1"}: { + header: header{ + "Content-Type": "text/event-stream", + }, + status: http.StatusOK, + done: allDone, // hang indefinitely + }, + {"POST", "123", notificationCancelled, ""}: {status: http.StatusAccepted}, + {"DELETE", "123", "", ""}: {optional: true}, + }, + } + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + defer close(allDone) // must be deferred *after* httpServer.Close, to avoid deadlock + + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + client := NewClient(testImpl, nil) + // Pin to 2025-11-25: the fixture's canned initialize response + // uses hardcoded id=1, which only matches when initialize is + // the first request. Under 2026-07-28 the client probes + // server/discover first. + cs, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() // ensure the session is closed, though we're also closing below + + // start the timer on the initial request + go func() { + <-time.After(2 * tick) + close(initialRequestDone) + }() + + // start the timer on the call cancellation + timeoutCtx, cancel := context.WithTimeout(ctx, test.cancelAfter) + defer cancel() + + go func() { + <-timeoutCtx.Done() + }() + + if _, err := cs.CallTool(timeoutCtx, &CallToolParams{ + Name: "tool", + }); err == nil { + t.Errorf("CallTool succeeded unexpectedly") + } + + // ...but cancellation should not break the session. + // Check that an arbitrary request succeeds. + if _, err := cs.ListTools(ctx, nil); err != nil { + t.Errorf("ListTools failed after cancellation") + } + }) + } +} + +// TestStreamableClientTransientErrors verifies that transient errors (timeouts, +// 5xx HTTP status codes) do not permanently break the client connection. +// This tests the fix for issue #683. +func TestStreamableClientTransientErrors(t *testing.T) { + ctx := context.Background() + + tests := []struct { + transientStatus int // HTTP status to return for the transient call + wantCallError bool // whether the transient call should error + wantSessionBroken bool // whether the session should be broken after + wantErrorContains string // substring expected in error message + }{ + { + transientStatus: http.StatusServiceUnavailable, + wantCallError: true, + wantSessionBroken: false, + wantErrorContains: "Service Unavailable", + }, + { + transientStatus: http.StatusBadGateway, + wantCallError: true, + wantSessionBroken: false, + wantErrorContains: "Bad Gateway", + }, + { + transientStatus: http.StatusGatewayTimeout, + wantCallError: true, + wantSessionBroken: false, + wantErrorContains: "Gateway Timeout", + }, + { + transientStatus: http.StatusTooManyRequests, + wantCallError: true, + wantSessionBroken: false, + wantErrorContains: "Too Many Requests", + }, + { + transientStatus: http.StatusUnauthorized, + wantCallError: true, + wantSessionBroken: true, + wantErrorContains: "Unauthorized", + }, + { + transientStatus: http.StatusNotFound, + wantCallError: true, + wantSessionBroken: true, + wantErrorContains: "not found", // NotFound has special handling + }, + } + + for _, test := range tests { + t.Run(http.StatusText(test.transientStatus), func(t *testing.T) { + var returnedError atomic.Bool + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + body: jsonBody(t, initResp), + }, + {"POST", "123", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + {"GET", "123", "", ""}: { + status: http.StatusMethodNotAllowed, + }, + {"POST", "123", methodListTools, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + responseFunc: func(r *jsonrpc.Request) (string, int) { + // First call returns transient error, subsequent calls succeed. + if !returnedError.Swap(true) && test.transientStatus != 0 { + return "", test.transientStatus + } + return jsonBody(t, resp(r.ID.Raw().(int64), &ListToolsResult{Tools: []*Tool{}}, nil)), 0 + }, + optional: true, + }, + {"DELETE", "123", "", ""}: {optional: true}, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + client := NewClient(testImpl, nil) + // Pin to 2025-11-25: the fixture's canned initialize response + // uses hardcoded id=1, which only matches when initialize is + // the first request. Under 2026-07-28 the client probes + // server/discover first. + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("Connect failed: %v", err) + } + defer session.Close() + + // First call: should trigger transient error. + _, err = session.ListTools(ctx, nil) + if test.wantCallError { + if err == nil { + t.Error("ListTools succeeded unexpectedly, want error") + } else if test.wantErrorContains != "" && !strings.Contains(err.Error(), test.wantErrorContains) { + t.Errorf("ListTools error = %q, want containing %q", err.Error(), test.wantErrorContains) + } + } else if err != nil { + t.Errorf("ListTools failed unexpectedly: %v", err) + } + + // Second call: verifies whether the session is still usable. + _, err = session.ListTools(ctx, nil) + if test.wantSessionBroken { + if err == nil { + t.Error("second ListTools succeeded unexpectedly, want session broken") + } + } else { + if err != nil { + t.Errorf("second ListTools failed unexpectedly: %v (session should survive transient errors)", err) + } + } + }) + } +} + +// TestStreamableClientRetryWithoutProgress verifies that the client fails after +// exceeding the retry limit when no progress is made (Last-Event-ID does not advance). +// This tests the fix for issue #679. +func TestStreamableClientRetryWithoutProgress(t *testing.T) { + // Speed up reconnection delays for testing. + const tick = 10 * time.Millisecond + defer func(delay int64) { + reconnectInitialDelay.Store(delay) + }(reconnectInitialDelay.Load()) + reconnectInitialDelay.Store(int64(tick)) + + // Use the fakeStreamableServer pattern like other tests to avoid race conditions. + ctx := context.Background() + const maxRetries = 2 + var retryCount atomic.Int32 + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "test-session", + }, + body: jsonBody(t, initResp), + }, + {"POST", "test-session", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + {"GET", "test-session", "", ""}: { + // Disable standalone SSE stream to simplify the test. + status: http.StatusMethodNotAllowed, + }, + {"POST", "test-session", methodCallTool, ""}: { + header: header{ + "Content-Type": "text/event-stream", + }, + // Return SSE stream with fixed event ID. + body: `id: fixed_1 +data: {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"test"}} + +`, + }, + // Resumption attempts with the same event ID (no progress). + {"GET", "test-session", "", "fixed_1"}: { + header: header{ + "Content-Type": "text/event-stream", + }, + responseFunc: func(r *jsonrpc.Request) (string, int) { + retryCount.Add(1) + // Return the same event ID - no progress. + return `id: fixed_1 +data: {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"retry"}} + +`, http.StatusOK + }, + }, + {"DELETE", "test-session", "", ""}: {optional: true}, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + MaxRetries: maxRetries, + } + client := NewClient(testImpl, nil) + // Pin to 2025-11-25: the fixture's canned initialize response uses + // hardcoded id=1, which only matches when initialize is the first + // request. Under 2026-07-28 the client probes server/discover first. + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("Connect failed: %v", err) + } + defer session.Close() + + // Make a call that will trigger reconnections without progress. + _, err = session.CallTool(ctx, &CallToolParams{Name: "test"}) + if err == nil { + t.Fatal("CallTool succeeded unexpectedly, want error due to exceeded retries") + } + + // Check that the error mentions exceeding retries without progress. + wantErr := "exceeded" + if !strings.Contains(err.Error(), wantErr) { + t.Errorf("CallTool error = %q, want containing %q", err.Error(), wantErr) + } + + // Verify that we actually retried the expected number of times. + // We expect maxRetries+1 attempts because we increment before checking the limit. + if got := retryCount.Load(); got != int32(maxRetries+1) { + t.Errorf("retry count = %d, want exactly %d", got, maxRetries+1) + } +} + +func TestStreamableClientDisableStandaloneSSE(t *testing.T) { + ctx := context.Background() + + tests := []struct { + name string + disableStandaloneSSE bool + expectGETRequest bool + }{ + { + name: "default behavior (standalone SSE enabled)", + disableStandaloneSSE: false, + expectGETRequest: true, + }, + { + name: "standalone SSE disabled", + disableStandaloneSSE: true, + expectGETRequest: false, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + getRequestKey := streamableRequestKey{"GET", "123", "", ""} + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + body: jsonBody(t, initResp), + }, + {"POST", "123", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + getRequestKey: { + header: header{ + "Content-Type": "text/event-stream", + }, + wantProtocolVersion: protocolVersion20251125, + optional: !test.expectGETRequest, + }, + {"DELETE", "123", "", ""}: { + optional: true, + }, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + DisableStandaloneSSE: test.disableStandaloneSSE, + } + client := NewClient(testImpl, nil) + // Pin to 2025-11-25: the fixture's canned initialize response + // uses hardcoded id=1, which only matches when initialize is + // the first request. Under 2026-07-28 the client probes + // server/discover first. + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("client.Connect() failed: %v", err) + } + + // Give some time for the standalone SSE connection to be established (if enabled) + time.Sleep(100 * time.Millisecond) + + // Verify the connection state + streamableConn, ok := session.mcpConn.(*streamableClientConn) + if !ok { + t.Fatalf("Expected *streamableClientConn, got %T", session.mcpConn) + } + + if got, want := streamableConn.disableStandaloneSSE, test.disableStandaloneSSE; got != want { + t.Errorf("disableStandaloneSSE field: got %v, want %v", got, want) + } + + // Clean up + if err := session.Close(); err != nil { + t.Errorf("closing session: %v", err) + } + + // Check if GET request was received + fake.calledMu.Lock() + getRequestReceived := false + if fake.called != nil { + getRequestReceived = fake.called[getRequestKey] + } + fake.calledMu.Unlock() + + if got, want := getRequestReceived, test.expectGETRequest; got != want { + t.Errorf("GET request received: got %v, want %v", got, want) + } + + // If we expected a GET request, verify it was actually received + if test.expectGETRequest { + if missing := fake.missingRequests(); len(missing) > 0 { + // Filter out optional requests + var requiredMissing []streamableRequestKey + for _, key := range missing { + if resp, ok := fake.responses[key]; ok && !resp.optional { + requiredMissing = append(requiredMissing, key) + } + } + if len(requiredMissing) > 0 { + t.Errorf("did not receive expected requests: %v", requiredMissing) + } + } + } else { + // If we didn't expect a GET request, verify it wasn't sent + if getRequestReceived { + t.Error("GET request was sent unexpectedly when DisableStandaloneSSE is true") + } + } + }) + } +} + +type mockOAuthHandler struct { + token *oauth2.Token + authorizeErr error + authorizeCalled bool +} + +func (h *mockOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) { + if h.token == nil { + return nil, nil + } + return oauth2.StaticTokenSource(h.token), nil +} + +func (h *mockOAuthHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error { + h.authorizeCalled = true + return h.authorizeErr +} + +func TestStreamableClientOAuth_AuthorizationHeader(t *testing.T) { + ctx := context.Background() + token := &oauth2.Token{AccessToken: "test-token"} + oauthHandler := &mockOAuthHandler{token: token} + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + body: jsonBody(t, initResp), + }, + {"POST", "123", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + {"GET", "123", "", ""}: { + header: header{ + "Content-Type": "text/event-stream", + }, + }, + {"DELETE", "123", "", ""}: {}, + }, + } + verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) { + if token != "test-token" { + return nil, auth.ErrInvalidToken + } + return &auth.TokenInfo{Expiration: time.Now().Add(time.Hour)}, nil + } + httpServer := httptest.NewServer(auth.RequireBearerToken(verifier, nil)(fake)) + t.Cleanup(httpServer.Close) + + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + OAuthHandler: oauthHandler, + } + client := NewClient(testImpl, nil) + // Pin to 2025-11-25: the fixture's canned initialize response uses + // hardcoded id=1, which only matches when initialize is the first + // request. Under 2026-07-28 the client probes server/discover first. + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("client.Connect() failed: %v", err) + } + session.Close() +} + +func TestStreamableClientOAuth_401(t *testing.T) { + ctx := context.Background() + oauthHandler := &mockOAuthHandler{token: nil} + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + body: jsonBody(t, initResp), + }, + }, + } + verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) { + // Accept any token. + return &auth.TokenInfo{Expiration: time.Now().Add(time.Hour)}, nil + } + httpServer := httptest.NewServer(auth.RequireBearerToken(verifier, nil)(fake)) + t.Cleanup(httpServer.Close) + + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + OAuthHandler: oauthHandler, + } + client := NewClient(testImpl, nil) + _, err := client.Connect(ctx, transport, nil) + if err == nil || !strings.Contains(err.Error(), "Unauthorized") { + t.Fatalf("client.Connect() error does not contain 'Unauthorized': %v", err) + } + + if !oauthHandler.authorizeCalled { + t.Errorf("expected Authorize to be called") + } +} + +// blockingCountingOAuthHandler is an OAuthHandler that blocks inside +// Authorize until the caller's context is cancelled, then returns a custom +// error that does NOT wrap context.Canceled. This mirrors real-world OAuth +// handlers that catch the cancellation internally and surface their own +// error type. The fix for #882 checks ctx.Err() directly rather than +// relying on the error from Authorize, so this must still trigger c.fail(). +// It records how many times Authorize is invoked. +type blockingCountingOAuthHandler struct { + mu sync.Mutex + callCount int +} + +func (h *blockingCountingOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) { + return nil, nil +} + +func (h *blockingCountingOAuthHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error { + h.mu.Lock() + h.callCount++ + h.mu.Unlock() + // Block until the caller's context is cancelled, mirroring an + // interactive OAuth flow that the user has abandoned. + <-ctx.Done() + // Return a custom error that does not wrap context.Canceled, as a + // real-world handler might. The code under test must check ctx.Err() + // to detect the cancellation, not this error. + return fmt.Errorf("oauth flow interrupted") +} + +func (h *blockingCountingOAuthHandler) Calls() int { + h.mu.Lock() + defer h.mu.Unlock() + return h.callCount +} + +// TestStreamableClientOAuth_CancelledAuthorize_NoReprompt is a regression +// test for #882. When OAuthHandler.Authorize returns a context-cancelled +// error, the connection must enter a failed state so that the cancellation +// notification the call layer sends in response to ctx cancellation does +// not flow back through the same broken auth path and re-invoke Authorize. +func TestStreamableClientOAuth_CancelledAuthorize_NoReprompt(t *testing.T) { + handler := &blockingCountingOAuthHandler{} + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + body: jsonBody(t, initResp), + }, + }, + } + verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) { + return &auth.TokenInfo{Expiration: time.Now().Add(time.Hour)}, nil + } + httpServer := httptest.NewServer(auth.RequireBearerToken(verifier, nil)(fake)) + t.Cleanup(httpServer.Close) + + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + OAuthHandler: handler, + } + client := NewClient(testImpl, nil) + + // Use a context with a tight deadline so the cancellation path runs + // while the auth flow is in progress. + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + + _, err := client.Connect(ctx, transport, nil) + if err == nil { + t.Fatal("expected client.Connect to fail") + } + + // Give the cancellation Notify path a moment to (try to) run. + time.Sleep(50 * time.Millisecond) + + // Authorize should be invoked exactly once. The bug in #882 caused + // it to be invoked a second time when the call layer sent the + // cancellation notification through the same auth-broken connection. + if got := handler.Calls(); got != 1 { + t.Errorf("expected Authorize to be called exactly 1 time, got %d", got) + } +} + +func TestTokenInfo(t *testing.T) { + ctx := context.Background() + + // Create a server with a tool that returns TokenInfo. + tokenInfo := func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, any, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: fmt.Sprintf("%v", req.Extra.TokenInfo)}}}, nil, nil + } + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "tokenInfo", Description: "return token info"}, tokenInfo) + + streamHandler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil) + verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) { + if token != "test-token" { + return nil, auth.ErrInvalidToken + } + return &auth.TokenInfo{ + Scopes: []string{"scope"}, + // Expiration is far, far in the future. + Expiration: time.Date(5000, 1, 2, 3, 4, 5, 0, time.UTC), + }, nil + } + handler := auth.RequireBearerToken(verifier, nil)(streamHandler) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + OAuthHandler: &mockOAuthHandler{token: &oauth2.Token{AccessToken: "test-token"}}, + } + client := NewClient(testImpl, nil) + session, err := client.Connect(ctx, transport, nil) + if err != nil { + t.Fatalf("client.Connect() failed: %v", err) + } + defer session.Close() + + res, err := session.CallTool(ctx, &CallToolParams{Name: "tokenInfo"}) + if err != nil { + t.Fatal(err) + } + if len(res.Content) == 0 { + t.Fatal("missing content") + } + tc, ok := res.Content[0].(*TextContent) + if !ok { + t.Fatal("not TextContent") + } + if g, w := tc.Text, "&{[scope] 5000-01-02 03:04:05 +0000 UTC map[]}"; g != w { + t.Errorf("got %q, want %q", g, w) + } +} + +// errTestAuthorizeFailed is a sentinel error returned by +// retrieveErrorOAuthHandler.Authorize(). +var errTestAuthorizeFailed = errors.New("authorize intentionally failed for test") + +// retrieveErrorOAuthHandler is a mock OAuthHandler that always returns +// an oauth2.RetrieveError from its TokenSource's Token() method. +type retrieveErrorOAuthHandler struct{} + +func (h *retrieveErrorOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) { + return h, nil +} + +func (h *retrieveErrorOAuthHandler) Token() (*oauth2.Token, error) { + return nil, &oauth2.RetrieveError{ + Response: &http.Response{StatusCode: http.StatusBadRequest}, + Body: []byte("test retrieve error"), + ErrorCode: "invalid_grant", + } +} + +func (h *retrieveErrorOAuthHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error { + return errTestAuthorizeFailed +} + +// TestStreamableClientOAuth_RetrieveError verifies that an invalid_grant RetrieveError +// from the OAuth token source correctly skips sending Authorization header and relies on +// the server's 401 response to trigger the Authorize fallback flow. +func TestStreamableClientOAuth_RetrieveError(t *testing.T) { + ctx := context.Background() + oauthHandler := &retrieveErrorOAuthHandler{} + + // Mock MCP server returns 401 Unauthorized to simulate a server rejecting + // the request that omitted the Authorization header. + httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + t.Cleanup(httpServer.Close) + + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + OAuthHandler: oauthHandler, + } + client := NewClient(testImpl, nil) + + // Attempt to connect. The Connect call will trigger the initialization request, + // which will fail to retrieve the token and proceed without auth header, receive 401, + // and invoke Authorize(). + _, err := client.Connect(ctx, transport, nil) + + // Expect the connection to fail with the sentinel error, not the RetrieveError. + if !errors.Is(err, errTestAuthorizeFailed) { + t.Fatalf("client.Connect() error = %v, want %v", err, errTestAuthorizeFailed) + } +} + +// discoverResult is the canned successful DiscoverResult returned by +// fakeStreamableServer setups in the tests below. +var discoverResult = &DiscoverResult{ + Meta: Meta{ + MetaKeyServerInfo: &Implementation{Name: "discoverServer", Version: "v1.0.0"}, + }, + SupportedVersions: []string{protocolVersion20260728}, + Capabilities: &ServerCapabilities{ + Tools: &ToolCapabilities{ListChanged: true}, + }, + Instructions: "test discover", +} + +// TestStreamableClientConnect_DiscoverSuccess verifies that Client.Connect on +// the streamable transport: +// - sends a POST server/discover with Mcp-Protocol-Version: 2026-07-28 and +// the SEP-2575 per-request _meta triple in the body, +// - on a successful DiscoverResult, skips the legacy initialize handshake +// entirely, and +// - seeds ClientSession.InitializeResult() from the discover response, +// picking a mutually-supported protocol version. +func TestStreamableClientConnect_DiscoverSuccess(t *testing.T) { + ctx := context.Background() + + var ( + gotDiscoverMu sync.Mutex + gotDiscover *jsonrpc.Request + ) + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodDiscover, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "sess-1", + }, + wantProtocolVersion: protocolVersion20260728, + responseFunc: func(r *jsonrpc.Request) (string, int) { + gotDiscoverMu.Lock() + gotDiscover = r + gotDiscoverMu.Unlock() + return jsonBody(t, resp(1, discoverResult, nil)), http.StatusOK + }, + }, + {"DELETE", "sess-1", "", ""}: {optional: true}, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + client := NewClient(testImpl, nil) + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer session.Close() + + if missing := fake.missingRequests(); len(missing) > 0 { + t.Errorf("missing expected requests: %v", missing) + } + + gotDiscoverMu.Lock() + defer gotDiscoverMu.Unlock() + if gotDiscover == nil { + t.Fatal("server did not receive server/discover") + } + + // Inspect the discover request body for the SEP-2575 _meta triple. + var body struct { + Meta map[string]any `json:"_meta"` + } + if err := json.Unmarshal(gotDiscover.Params, &body); err != nil { + t.Fatalf("decoding discover params: %v", err) + } + if v, _ := body.Meta[MetaKeyProtocolVersion].(string); v != protocolVersion20260728 { + t.Errorf("_meta[%s] = %q, want %q", MetaKeyProtocolVersion, v, protocolVersion20260728) + } + if _, ok := body.Meta[MetaKeyClientInfo]; !ok { + t.Errorf("_meta[%s] missing", MetaKeyClientInfo) + } + if _, ok := body.Meta[MetaKeyClientCapabilities]; !ok { + t.Errorf("_meta[%s] missing", MetaKeyClientCapabilities) + } + + ir := session.InitializeResult() + if ir == nil { + t.Fatal("InitializeResult is nil after Connect") + } + if got, want := ir.ProtocolVersion, protocolVersion20260728; got != want { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q", got, want) + } + if ir.ServerInfo == nil || ir.ServerInfo.Name != "discoverServer" { + t.Errorf("InitializeResult.ServerInfo = %+v, want Name=discoverServer", ir.ServerInfo) + } + if ir.Instructions != "test discover" { + t.Errorf("InitializeResult.Instructions = %q, want %q", ir.Instructions, "test discover") + } +} + +// TestStreamableClientConnSetMCPHeaders_ProtocolVersion covers +// streamableClientConn.setMCPHeaders' selection of the Mcp-Protocol-Version +// header value. +// +// Ordinarily initializedResult is populated by sessionUpdated, called +// through a type assertion to the unexported clientConnection interface +// (see Client.Connect). That assertion silently fails, leaving +// initializedResult nil for the life of the session, whenever the +// Connection returned by a Transport is wrapped by another type exposing +// only the exported Connection interface (a real pattern for transports +// that intercept traffic, e.g. to filter notifications): Go does not +// promote unexported interface methods across an embedded interface +// boundary. Every SEP-2575 (>= 2026-07-28) request already carries its own +// `_meta.protocolVersion` field, so setMCPHeaders falls back to reading it +// from the outgoing message when initializedResult is unset. +func TestStreamableClientConnSetMCPHeaders_ProtocolVersion(t *testing.T) { + tests := []struct { + name string + initializedResult *InitializeResult + msg jsonrpc.Message + ctxVersion string + want string + }{ + { + name: "nil checked", + initializedResult: nil, + msg: (*jsonrpc.Request)(nil), + want: "", + }, + { + name: "message meta wins when initializedResult unset", + initializedResult: nil, + msg: req(1, methodListTools, &ListToolsParams{Meta: Meta{MetaKeyProtocolVersion: protocolVersion20260728}}), + want: protocolVersion20260728, + }, + { + name: "initializedResult used when message has no meta", + initializedResult: &InitializeResult{ProtocolVersion: protocolVersion20251125}, + msg: req(1, methodListTools, &ListToolsParams{}), + want: protocolVersion20251125, + }, + { + name: "initializedResult used for nil message (GET/DELETE)", + initializedResult: &InitializeResult{ProtocolVersion: protocolVersion20251125}, + msg: nil, + want: protocolVersion20251125, + }, + { + name: "message meta preferred over stale initializedResult", + initializedResult: &InitializeResult{ProtocolVersion: protocolVersion20251125}, + msg: req(1, methodListTools, &ListToolsParams{Meta: Meta{MetaKeyProtocolVersion: protocolVersion20260728}}), + want: protocolVersion20260728, + }, + { + name: "no header when neither source is set", + initializedResult: nil, + msg: req(1, methodListTools, &ListToolsParams{}), + want: "", + }, + { + // A process that is both a server and a client (a proxy) carries the + // inbound request's version on the context of the outgoing call. The + // version this connection negotiated must win over it. + name: "initializedResult preferred over request context", + initializedResult: &InitializeResult{ProtocolVersion: protocolVersion20251125}, + msg: req(1, methodListTools, &ListToolsParams{}), + ctxVersion: protocolVersion20260728, + want: protocolVersion20251125, + }, + { + name: "request context used when initializedResult unset", + initializedResult: nil, + msg: req(1, methodListTools, &ListToolsParams{}), + ctxVersion: protocolVersion20260728, + want: protocolVersion20260728, + }, + { + name: "message meta preferred over request context", + initializedResult: nil, + msg: req(1, methodListTools, &ListToolsParams{Meta: Meta{MetaKeyProtocolVersion: protocolVersion20251125}}), + ctxVersion: protocolVersion20260728, + want: protocolVersion20251125, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + conn := &streamableClientConn{initializedResult: tt.initializedResult} + httpReq, err := http.NewRequest(http.MethodPost, "http://test.invalid", nil) + if err != nil { + t.Fatal(err) + } + if tt.ctxVersion != "" { + httpReq = httpReq.WithContext(context.WithValue(httpReq.Context(), protocolVersionContextKey{}, tt.ctxVersion)) + } + if err := conn.setMCPHeaders(httpReq, tt.msg); err != nil { + t.Fatalf("setMCPHeaders: %v", err) + } + if got := httpReq.Header.Get(protocolVersionHeader); got != tt.want { + t.Errorf("Mcp-Protocol-Version header = %q, want %q", got, tt.want) + } + }) + } +} + +// TestStreamableClientConnect_DiscoverMethodNotFound verifies that Client.Connect +// falls back to the legacy initialize handshake when the server responds to +// server/discover with a JSON-RPC "Method not found" error. +func TestStreamableClientConnect_DiscoverMethodNotFound(t *testing.T) { + ctx := context.Background() + + // Each request gets a fresh jsonrpc2 ID from the same client connection. + // Use responseFunc to echo the request's ID back so the client matches + // the response to the in-flight call regardless of ordering. + echoErr := func(err error) func(*jsonrpc.Request) (string, int) { + return func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, &jsonrpc.Response{ID: r.ID, Error: err.(*jsonrpc.Error)}), http.StatusOK + } + } + echoResult := func(result any) func(*jsonrpc.Request) (string, int) { + return func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(result)}), http.StatusOK + } + } + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodDiscover, ""}: { + header: header{"Content-Type": "application/json"}, + wantProtocolVersion: protocolVersion20260728, + responseFunc: echoErr(&jsonrpc.Error{ + Code: jsonrpc.CodeMethodNotFound, + Message: "method not found", + }), + }, + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "fallback", + }, + responseFunc: echoResult(initResult), + }, + {"POST", "fallback", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + {"GET", "fallback", "", ""}: { + header: header{"Content-Type": "text/event-stream"}, + wantProtocolVersion: protocolVersion20251125, + optional: true, + }, + {"DELETE", "fallback", "", ""}: {optional: true}, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + client := NewClient(testImpl, nil) + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer session.Close() + + // The fallback initialize explicitly requests protocolVersion20251125 + // (see client.go), so the server negotiates that version. + if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20251125 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (initialize fallback)", got, protocolVersion20251125) + } +} + +// TestStreamableClientConnect_DiscoverUnsupportedVersion verifies that +// Client.Connect falls back to the legacy initialize handshake when the +// server responds to server/discover with the SEP-2575 +// UnsupportedProtocolVersionError JSON-RPC code. +func TestStreamableClientConnect_DiscoverUnsupportedVersion(t *testing.T) { + ctx := context.Background() + + echoErr := func(err error) func(*jsonrpc.Request) (string, int) { + return func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, &jsonrpc.Response{ID: r.ID, Error: err.(*jsonrpc.Error)}), http.StatusOK + } + } + echoResult := func(result any) func(*jsonrpc.Request) (string, int) { + return func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(result)}), http.StatusOK + } + } + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodDiscover, ""}: { + header: header{"Content-Type": "application/json"}, + wantProtocolVersion: protocolVersion20260728, + responseFunc: echoErr(&jsonrpc.Error{ + Code: CodeUnsupportedProtocolVersion, + Message: "unsupported protocol version", + }), + }, + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "fallback", + }, + responseFunc: echoResult(initResult), + }, + {"POST", "fallback", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + {"GET", "fallback", "", ""}: { + header: header{"Content-Type": "text/event-stream"}, + wantProtocolVersion: protocolVersion20251125, + optional: true, + }, + {"DELETE", "fallback", "", ""}: {optional: true}, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + client := NewClient(testImpl, nil) + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer session.Close() + + // The fallback initialize explicitly requests protocolVersion20251125 + // (see client.go), so the server negotiates that version. + if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20251125 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (initialize fallback)", got, protocolVersion20251125) + } +} + +// TestStreamableClientConnect_DiscoverMethodNotFoundVPre verifies that +// Client.Connect falls back to the legacy initialize handshake when a +// pre-SEP-2575 (vPre) server rejects server/discover. +func TestStreamableClientConnect_DiscoverMethodNotFoundVPre(t *testing.T) { + ctx := context.Background() + + echoResult := func(result any) func(*jsonrpc.Request) (string, int) { + return func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(result)}), http.StatusOK + } + } + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodDiscover, ""}: { + wantProtocolVersion: protocolVersion20260728, + // Reproduce the exact body a vPre server produces via + // http.Error(w, err.Error(), 400) where err comes from + // checkRequest. http.Error appends a trailing newline. + body: "JSON RPC not handled: \"server/discover\" unsupported\n", + status: http.StatusBadRequest, + header: header{"Content-Type": "text/plain; charset=utf-8"}, + }, + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "fallback", + }, + responseFunc: echoResult(initResult), + }, + {"POST", "fallback", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + {"DELETE", "fallback", "", ""}: {optional: true}, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + DisableStandaloneSSE: true, + } + client := NewClient(testImpl, nil) + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer session.Close() + + // The fallback initialize explicitly requests protocolVersion20251125 + // (see client.go), so the server negotiates that version. + if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20251125 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (initialize fallback)", got, protocolVersion20251125) + } +} + +// TestStreamableClientConnect_DiscoverUnsupportedProtocolVersion verifies that +// Client.Connect falls back to the legacy initialize handshake when a +// server rejects server/discover with a plain HTTP 400 containing "Unsupported protocol version". +func TestStreamableClientConnect_DiscoverUnsupportedVersionVPre(t *testing.T) { + ctx := context.Background() + + echoResult := func(result any) func(*jsonrpc.Request) (string, int) { + return func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(result)}), http.StatusOK + } + } + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodDiscover, ""}: { + wantProtocolVersion: protocolVersion20260728, + body: "Bad Request: Unsupported protocol version\n", + status: http.StatusBadRequest, + header: header{"Content-Type": "text/plain; charset=utf-8"}, + }, + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "fallback", + }, + responseFunc: echoResult(initResult), + }, + {"POST", "fallback", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + {"DELETE", "fallback", "", ""}: {optional: true}, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + DisableStandaloneSSE: true, + } + client := NewClient(testImpl, nil) + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer session.Close() + + // The fallback initialize explicitly requests protocolVersion20251125 + // (see client.go), so the server negotiates that version. + if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20251125 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (initialize fallback)", got, protocolVersion20251125) + } +} + +// TestStreamableClientConnect_DiscoverUnsupportedVersionNegotiation verifies that +// when Client.Connect over a streamable transport receives an +// UnsupportedProtocolVersion error containing Data.Supported, it negotiates a +// mutually supported version and retries server/discover. +func TestStreamableClientConnect_DiscoverUnsupportedVersionNegotiation(t *testing.T) { + ctx := context.Background() + + const unsupportedClientVersion = "2099-12-31" + + var discoverCalls atomic.Int32 + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodDiscover, ""}: { + header: header{ + "Content-Type": "application/json", + }, + responseFunc: func(r *jsonrpc.Request) (string, int) { + n := discoverCalls.Add(1) + if n == 1 { + data, _ := json.Marshal(UnsupportedProtocolVersionData{ + Supported: []string{protocolVersion20260728}, + }) + respMsg := &jsonrpc.Response{ + ID: r.ID, + Error: &jsonrpc.Error{ + Code: CodeUnsupportedProtocolVersion, + Message: "unsupported protocol version", + Data: data, + }, + } + return jsonBody(t, respMsg), http.StatusOK + } + return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(discoverResult)}), http.StatusOK + }, + }, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + client := NewClient(testImpl, nil) + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: unsupportedClientVersion}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer session.Close() + + if got, want := discoverCalls.Load(), int32(2); got != want { + t.Errorf("discover call count = %d, want %d", got, want) + } + if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20260728 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q", got, protocolVersion20260728) + } +} + +// TestStreamableClientHandlerErrorPropagation verifies that per-call +// handler-level HTTP errors carrying a JSON-RPC error body do not tear down +// the session, and that setting MCPGODEBUG=noprotocolerrorbody=1 restores the +// pre-fix behavior in which non-transient errors permanently failed the +// connection. +func TestStreamableClientHandlerErrorPropagation(t *testing.T) { + ctx := context.Background() + + // Build a JSON-RPC error response body for a given tools/call request. + jsonRPCErrorBody := func(id int64, code int64, msg string) string { + return jsonBody(t, resp(id, nil, &jsonrpc.Error{Code: code, Message: msg})) + } + + tests := []struct { + name string + callStatus int // HTTP status returned for the tools/call + callBody string // response body for the tools/call ("" for empty; may be a JSON-RPC error) + disableBodyDecode bool // if true, sets noprotocolerrorbody="1" for the duration of the test + wantErrRejected bool // whether the returned error should match errors.Is(err, jsonrpc2.ErrRejected) + wantSessionAlive bool // whether a subsequent call on the same session should succeed + }{ + { + // SEP-2575 §"Missing Required Capabilities": server returns + // -32021 at HTTP 400. Session must survive. + name: "400 with JSON-RPC error body (default)", + callStatus: http.StatusBadRequest, + callBody: jsonRPCErrorBody(1, CodeMissingRequiredClientCapabilities, "missing capability"), + wantErrRejected: true, + wantSessionAlive: true, + }, + { + // Same server response, but the user opted out via + // MCPGODEBUG=noprotocolerrorbody=1. Pre-fix behavior: the + // session is torn down. + name: "400 with JSON-RPC error body (noprotocolerrorbody=1)", + callStatus: http.StatusBadRequest, + callBody: jsonRPCErrorBody(1, CodeMissingRequiredClientCapabilities, "missing capability"), + disableBodyDecode: true, + wantErrRejected: false, + wantSessionAlive: false, + }, + { + // Legacy server returns plain-text 400 with no JSON-RPC + // body: cannot be classified as a per-call rejection, so the + // session is still torn down. + name: "400 with plain-text body", + callStatus: http.StatusBadRequest, + callBody: "Bad Request", + wantErrRejected: false, + wantSessionAlive: false, + }, + { + // SEP-2575: server returns -32601 for an unimplemented + // method at HTTP 404. Must be treated as a per-call + // rejection, not a terminated session. + name: "404 with JSON-RPC MethodNotFound body", + callStatus: http.StatusNotFound, + callBody: jsonRPCErrorBody(1, jsonrpc.CodeMethodNotFound, `method not found: "tools/call"`), + wantErrRejected: true, + wantSessionAlive: true, + }, + { + // A bare 404 with no JSON-RPC body means the session has + // been terminated. Session must not survive. + name: "404 with empty body (terminated session)", + callStatus: http.StatusNotFound, + callBody: "", + wantErrRejected: false, + wantSessionAlive: false, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if test.disableBodyDecode { + prev := noprotocolerrorbody + noprotocolerrorbody = "1" + t.Cleanup(func() { noprotocolerrorbody = prev }) + } + + // Track how many tools/call requests we've served: first + // returns the error under test, subsequent calls succeed so + // we can observe whether the session is still alive. + var callsServed atomic.Int32 + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodDiscover, ""}: { + header: header{ + "Content-Type": "application/json", + }, + wantProtocolVersion: protocolVersion20260728, + responseFunc: func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, resp(r.ID.Raw().(int64), discoverResult, nil)), http.StatusOK + }, + }, + {"POST", "", "tools/call", ""}: { + header: header{"Content-Type": "application/json"}, + responseFunc: func(r *jsonrpc.Request) (string, int) { + if callsServed.Add(1) == 1 { + return test.callBody, test.callStatus + } + return jsonBody(t, resp(r.ID.Raw().(int64), &CallToolResult{}, nil)), 0 + }, + optional: true, + }, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + client := NewClient(testImpl, nil) + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer session.Close() + + // First call: triggers the error under test. + _, err = session.CallTool(ctx, &CallToolParams{Name: "nonexistent"}) + if err == nil { + t.Fatal("first CallTool succeeded unexpectedly, want error") + } + if got := errors.Is(err, jsonrpc2.ErrRejected); got != test.wantErrRejected { + t.Errorf("errors.Is(err, ErrRejected) = %v, want %v (err = %v)", got, test.wantErrRejected, err) + } + + // Second call: verifies whether the session survived. + _, err = session.CallTool(ctx, &CallToolParams{Name: "nonexistent"}) + if test.wantSessionAlive { + if err != nil { + t.Errorf("second CallTool failed: %v (session should survive per-call rejection)", err) + } + } else { + if err == nil { + t.Error("second CallTool succeeded unexpectedly, want session torn down") + } + } + }) + } +} + +// TestStreamableClient_StatelessSubscriptionsListen404 verifies that in stateless mode +// (protocol 2026-07-28), when a server rejects subscriptions/listen with 404 (e.g. Copilot MCP server), +// the client connection survives and subsequent RPCs (such as tools/list) succeed. +func TestStreamableClient_StatelessSubscriptionsListen404(t *testing.T) { + ctx := t.Context() + + var listenServed atomic.Bool + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodDiscover, ""}: { + header: header{ + "Content-Type": "application/json", + }, + wantProtocolVersion: protocolVersion20260728, + responseFunc: func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, resp(r.ID.Raw().(int64), discoverResult, nil)), http.StatusOK + }, + }, + {"POST", "", methodSubscriptionsListen, ""}: { + header: header{"Content-Type": "text/plain"}, + responseFunc: func(r *jsonrpc.Request) (string, int) { + listenServed.Store(true) + return "404 Not Found", http.StatusNotFound + }, + optional: true, + }, + {"POST", "", methodListTools, ""}: { + header: header{"Content-Type": "application/json"}, + responseFunc: func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, resp(r.ID.Raw().(int64), &ListToolsResult{Tools: []*Tool{}}, nil)), http.StatusOK + }, + optional: true, + }, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + client := NewClient(testImpl, &ClientOptions{ + ToolListChangedHandler: func(ctx context.Context, req *ToolListChangedRequest) {}, + }) + + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("Connect failed: %v", err) + } + defer session.Close() + + // tools/list must succeed even though subscriptions/listen returned 404. + res, err := session.ListTools(ctx, nil) + if err != nil { + t.Fatalf("ListTools failed: %v", err) + } + if res == nil { + t.Fatal("ListTools result is nil") + } + if !listenServed.Load() { + t.Fatal("subscriptions/listen was not called") + } +} diff --git a/mcp/streamable_example_test.go b/mcp/streamable_example_test.go index f1cdf90..66834f2 100644 --- a/mcp/streamable_example_test.go +++ b/mcp/streamable_example_test.go @@ -32,10 +32,10 @@ func ExampleStreamableHTTPHandler() { defer httpServer.Close() // The SDK is currently permissive of some missing keys in "params". - resp := mustPostMessage(`{"jsonrpc": "2.0", "id": 1, "method":"initialize", "params": {}}`, httpServer.URL) + resp := mustPostMessage(`{"jsonrpc": "2.0", "id": 1, "method":"initialize", "params": {"protocolVersion":"2025-11-25"}}`, httpServer.URL) fmt.Println(resp) // Output: - // {"jsonrpc":"2.0","id":1,"result":{"capabilities":{"logging":{}},"protocolVersion":"2025-06-18","serverInfo":{"name":"server","version":"v0.1.0"}}} + // {"jsonrpc":"2.0","id":1,"result":{"capabilities":{"logging":{}},"protocolVersion":"2025-11-25","serverInfo":{"name":"server","version":"v0.1.0"}}} } // !-streamablehandler diff --git a/mcp/streamable_headers.go b/mcp/streamable_headers.go new file mode 100644 index 0000000..c2b8dfd --- /dev/null +++ b/mcp/streamable_headers.go @@ -0,0 +1,545 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package mcp + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "log/slog" + "math" + "net/http" + "strconv" + "strings" + + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" +) + +const ( + protocolVersionHeader = "Mcp-Protocol-Version" + sessionIDHeader = "Mcp-Session-Id" + lastEventIDHeader = "Last-Event-ID" + methodHeader = "Mcp-Method" + nameHeader = "Mcp-Name" + paramHeaderPrefix = "Mcp-Param-" + minVersionForStandardHeaders = protocolVersion20260728 + base64Prefix = "=?base64?" + base64Suffix = "?=" +) + +func extractName(method string, params json.RawMessage) (string, bool) { + switch method { + case "tools/call": + var p CallToolParams + if err := internaljson.Unmarshal(params, &p); err == nil { + return p.Name, true + } + case "prompts/get": + var p GetPromptParams + if err := internaljson.Unmarshal(params, &p); err == nil { + return p.Name, true + } + case "resources/read": + var p ReadResourceParams + if err := internaljson.Unmarshal(params, &p); err == nil { + return p.URI, true + } + } + + return "", false +} + +// headerSchemaProperty captures the fields needed for x-mcp-header processing. +type headerSchemaProperty struct { + Type string `json:"type"` + XMCPHeader json.RawMessage `json:"x-mcp-header,omitempty"` + Properties map[string]headerSchemaProperty `json:"properties,omitempty"` +} + +// unmarshalSchemaProperties normalizes any InputSchema type +// (*jsonschema.Schema, map[string]any, or json.RawMessage) into a common +// representation by marshaling to JSON and unmarshaling only the fields we need. +func unmarshalSchemaProperties(schema any) map[string]headerSchemaProperty { + var s headerSchemaProperty + if err := remarshal(schema, &s); err != nil { + return nil + } + return s.Properties +} + +// paramHeaderBinding maps a (possibly nested) input-schema property to the +// HTTP header it carries. +type paramHeaderBinding struct { + Path []string + Header string +} + +// extractParamHeaderAnnotations returns the bindings for every property in +// the tool's InputSchema that has an x-mcp-header annotation +func extractParamHeaderAnnotations(tool *Tool) []paramHeaderBinding { + props := unmarshalSchemaProperties(tool.InputSchema) + if len(props) == 0 { + return nil + } + var result []paramHeaderBinding + result = collectParamHeaderAnnotations(props, nil, result) + if len(result) == 0 { + return nil + } + return result +} + +// collectParamHeaderAnnotations walks the schema properties and records every +// x-mcp-header annotation it finds, keyed by the property-name path. +func collectParamHeaderAnnotations(props map[string]headerSchemaProperty, prefix []string, out []paramHeaderBinding) []paramHeaderBinding { + for propName, prop := range props { + path := make([]string, len(prefix)+1) + copy(path, prefix) + path[len(prefix)] = propName + + var headerName string + if err := json.Unmarshal(prop.XMCPHeader, &headerName); err == nil && headerName != "" { + out = append(out, paramHeaderBinding{Path: path, Header: headerName}) + } + if len(prop.Properties) > 0 { + out = collectParamHeaderAnnotations(prop.Properties, path, out) + } + } + return out +} + +// lookupArgument navigates the arguments object using the given property-name +// path and returns the raw JSON value at that location. It reports whether +// the value was found. +func lookupArgument(args map[string]json.RawMessage, path []string) (json.RawMessage, bool) { + if len(path) == 0 { + return nil, false + } + cur, ok := args[path[0]] + if !ok { + return nil, false + } + for _, part := range path[1:] { + var obj map[string]json.RawMessage + if err := internaljson.Unmarshal(cur, &obj); err != nil { + return nil, false + } + cur, ok = obj[part] + if !ok { + return nil, false + } + } + return cur, true +} + +// maxSafeInteger and minSafeInteger bound the integer values that can be +// faithfully represented as IEEE-754 double-precision floats. +const ( + maxSafeInteger = 1<<53 - 1 // 2^53 - 1 = 9007199254740991 + minSafeInteger = -(1<<53 - 1) // -(2^53 - 1) = -9007199254740991 +) + +// unmarshalPrimitive unmarshals a JSON value into the Go representation used +// for x-mcp-header processing per SEP-2243: +// +// - JSON string -> string +// - JSON boolean -> bool +// - JSON integer (within the JavaScript safe-integer range) -> int64 +// +// JSON numbers that are non-integers (have a fractional part, NaN, or ±Inf) +// or integers outside the safe range are rejected because the `number` type +// is not permitted for x-mcp-header parameters; only integer, string, boolean +// are allowed. +func unmarshalPrimitive(raw json.RawMessage) any { + var val any + if err := internaljson.Unmarshal(raw, &val); err != nil { + return nil + } + switch v := val.(type) { + case string, bool: + return v + case float64: + if math.IsNaN(v) || math.IsInf(v, 0) || v != math.Trunc(v) { + return nil + } + if v < minSafeInteger || v > maxSafeInteger { + return nil + } + return int64(v) + default: + return nil + } +} + +// primitiveToString formats an x-mcp-header value (as produced by +// [unmarshalPrimitive]) to its canonical header string representation per +// SEP-2243. Returns false if value is not one of the permitted primitive +// types (string, bool, int64). +func primitiveToString(value any) (string, bool) { + switch v := value.(type) { + case string: + return v, true + case bool: + return fmt.Sprintf("%t", v), true + case int64: + return strconv.FormatInt(v, 10), true + default: + return "", false + } +} + +// setStandardHeaders populates standard MCP headers. +// It requires the protocol version header to be set. +func setStandardHeaders(ctx context.Context, header http.Header, msg jsonrpc.Message) { + if msg == nil { + return + } + if header.Get(protocolVersionHeader) == "" || header.Get(protocolVersionHeader) < minVersionForStandardHeaders { + return + } + + switch msg := msg.(type) { + case *jsonrpc.Request: + header.Set(methodHeader, msg.Method) + if name, ok := extractName(msg.Method, msg.Params); ok { + // Names are only SHOULD-constrained to header-safe characters, so a + // name outside the safe set is Base64-wrapped with the =?base64?...?= + // sentinel, the same rule applied to Mcp-Param-* values. + if encoded, ok := encodeHeaderValue(name); ok { + header.Set(nameHeader, encoded) + } + } + if msg.Method == "tools/call" { + if tool, ok := ctx.Value(toolContextKey).(*Tool); ok && tool != nil { + for k, v := range generateParamHeaders(tool, msg.Params) { + header.Set(k, v) + } + } + } + } +} + +// generateParamHeaders reads x-mcp-header annotations from the tool's InputSchema +// and returns the Mcp-Param-{Name} headers to be set on the HTTP request. +func generateParamHeaders(tool *Tool, params json.RawMessage) map[string]string { + paramHeaders := extractParamHeaderAnnotations(tool) + if len(paramHeaders) == 0 { + return nil + } + + var raw struct { + Arguments map[string]json.RawMessage `json:"arguments"` + } + if err := internaljson.Unmarshal(params, &raw); err != nil || raw.Arguments == nil { + return nil + } + + res := make(map[string]string) + for _, b := range paramHeaders { + argRaw, ok := lookupArgument(raw.Arguments, b.Path) + if !ok { + continue + } + if string(argRaw) == "null" { + continue + } + val := unmarshalPrimitive(argRaw) + if val == nil { + continue + } + encoded, ok := encodeHeaderValue(val) + if !ok { + continue + } + res[paramHeaderPrefix+b.Header] = encoded + } + return res +} + +// filterValidTools returns only tools that have valid +// x-mcp-header annotations. Invalid tools are logged and excluded. +func filterValidTools(logger *slog.Logger, tools []*Tool) []*Tool { + logger = ensureLogger(logger) + result := make([]*Tool, 0, len(tools)) + for _, tool := range tools { + // A nil tool (e.g. from a malformed "tools":[null] response) is + // invalid; exclude it instead of dereferencing it. + if tool == nil { + logger.Error("excluding nil tool from tools/list") + continue + } + if err := validateParamHeaderAnnotations(tool); err != nil { + logger.Error("excluding tool from tools/list", "tool", tool.Name, "error", err) + continue + } + result = append(result, tool) + } + return result +} + +// validateParamHeaderAnnotations checks that a tool's x-mcp-header annotations +// are valid. Annotations may appear on properties at any nesting +// depth within the inputSchema and must be unique across all of them. +func validateParamHeaderAnnotations(tool *Tool) error { + props := unmarshalSchemaProperties(tool.InputSchema) + if len(props) == 0 { + return nil + } + seen := make(map[string]bool) + return validateParamHeadersIn(props, "", seen) +} + +func validateParamHeadersIn(props map[string]headerSchemaProperty, prefix string, seen map[string]bool) error { + for propName, prop := range props { + path := propName + if prefix != "" { + path = prefix + "." + propName + } + if prop.XMCPHeader != nil { + if prop.Type != "string" && prop.Type != "integer" && prop.Type != "boolean" { + return fmt.Errorf("property %q: x-mcp-header can only be applied to primitive types (integer, string, boolean), got %q", path, prop.Type) + } + var headerName string + if err := json.Unmarshal(prop.XMCPHeader, &headerName); err != nil || headerName == "" { + return fmt.Errorf("property %q: x-mcp-header must be a non-empty string", path) + } + if err := validateHeaderName(headerName); err != nil { + return fmt.Errorf("property %q: %w", path, err) + } + lower := strings.ToLower(headerName) + if seen[lower] { + return fmt.Errorf("property %q: duplicate x-mcp-header value %q (case-insensitive)", path, headerName) + } + seen[lower] = true + } + if len(prop.Properties) > 0 { + if err := validateParamHeadersIn(prop.Properties, path, seen); err != nil { + return err + } + } + } + return nil +} + +// validateHeaderName checks that a header name matches the HTTP field-name +// token syntax (1*tchar). +func validateHeaderName(name string) error { + if name == "" { + return fmt.Errorf("x-mcp-header value must be a non-empty string") + } + for _, c := range name { + if !isTChar(c) { + return fmt.Errorf("x-mcp-header value %q contains invalid character %q", name, c) + } + } + return nil +} + +// isTChar reports whether c is a valid HTTP token character (tchar) +// +// tchar = "!" / "#" / "$" / "%" / "&" / "'" / "*" / "+" / "-" / "." / +// "^" / "_" / "`" / "|" / "~" / DIGIT / ALPHA +func isTChar(c rune) bool { + switch { + case c >= '0' && c <= '9': + return true + case c >= 'A' && c <= 'Z': + return true + case c >= 'a' && c <= 'z': + return true + } + switch c { + case '!', '#', '$', '%', '&', '\'', '*', '+', '-', '.', + '^', '_', '`', '|', '~': + return true + } + return false +} + +func validateMcpHeaders(header http.Header, msg jsonrpc.Message, toolLookup func(string) (*serverTool, bool)) error { + protocolVersion := header.Get(protocolVersionHeader) + if protocolVersion == "" || protocolVersion < minVersionForStandardHeaders { + return nil + } + + switch msg := msg.(type) { + case *jsonrpc.Request: + methodInHeader := header.Get(methodHeader) + if methodInHeader == "" { + return errors.New("missing required Mcp-Method header") + } + if methodInHeader != msg.Method { + return fmt.Errorf("header mismatch: Mcp-Method header value '%s' does not match body value '%s'", methodInHeader, msg.Method) + } + + var nameInBody string + if msg.Method == "tools/call" || msg.Method == "resources/read" || msg.Method == "prompts/get" { + nameInHeader := header.Get(nameHeader) + if nameInHeader == "" { + return fmt.Errorf("missing required Mcp-Name header for method %q", msg.Method) + } + // A name that is not header-safe may be Base64-wrapped by the client + // (=?base64?...?=), so decode before comparing, as with Mcp-Param-*. + decodedName, ok := decodeHeaderValue(nameInHeader) + if !ok { + return errors.New("header mismatch: Mcp-Name header contains invalid Base64 encoding") + } + nameInBody, ok = extractName(msg.Method, msg.Params) + if !ok { + return fmt.Errorf("failed to extract name from parameters for method %q", msg.Method) + } + if decodedName != nameInBody { + return fmt.Errorf("header mismatch: Mcp-Name header value '%s' does not match body value '%s'", decodedName, nameInBody) + } + } + + if msg.Method == "tools/call" && toolLookup != nil { + if st, ok := toolLookup(nameInBody); ok && st != nil { + if err := validateParamHeaders(header, msg, st.tool); err != nil { + return err + } + } + } + } + return nil +} + +func validateParamHeaders(header http.Header, msg *jsonrpc.Request, tool *Tool) error { + paramHeaders := extractParamHeaderAnnotations(tool) + if len(paramHeaders) == 0 { + return nil + } + + var raw struct { + Arguments map[string]json.RawMessage `json:"arguments"` + } + if err := internaljson.Unmarshal(msg.Params, &raw); err != nil { + return nil + } + + for _, b := range paramHeaders { + fullHeader := paramHeaderPrefix + b.Header + headerVal := header.Get(fullHeader) + argRaw, argExists := lookupArgument(raw.Arguments, b.Path) + + if !argExists || string(argRaw) == "null" { + if headerVal != "" { + return fmt.Errorf("header mismatch: unexpected %s header for absent or null parameter %q", fullHeader, strings.Join(b.Path, ".")) + } + continue + } + + if headerVal == "" { + return fmt.Errorf("header mismatch: missing %s header for parameter %q", fullHeader, strings.Join(b.Path, ".")) + } + + decoded, ok := decodeHeaderValue(headerVal) + if !ok { + return fmt.Errorf("header mismatch: %s header contains invalid Base64 encoding", fullHeader) + } + + bodyVal := unmarshalPrimitive(argRaw) + if bodyVal == nil { + return fmt.Errorf("header mismatch: %s header present but body parameter %q is not a primitive type", fullHeader, strings.Join(b.Path, ".")) + } + + if !primitiveEqual(decoded, bodyVal) { + return fmt.Errorf("header mismatch: %s header value '%s' does not match body value", fullHeader, headerVal) + } + } + return nil +} + +// primitiveEqual reports whether the (decoded) header string equals the +// JSON-derived body value. +func primitiveEqual(headerStr string, bodyVal any) bool { + if bodyInt, ok := bodyVal.(int64); ok { + headerNum, err := strconv.ParseFloat(headerStr, 64) + if err != nil { + return false + } + if math.IsNaN(headerNum) || math.IsInf(headerNum, 0) || headerNum != math.Trunc(headerNum) { + return false + } + if headerNum < minSafeInteger || headerNum > maxSafeInteger { + return false + } + return int64(headerNum) == bodyInt + } + expected, ok := primitiveToString(bodyVal) + if !ok { + return false + } + return headerStr == expected +} + +// encodeHeaderValue converts a parameter value to an HTTP header-safe string +// per the SEP-2243 encoding rules: +// - string: used as-is if safe ASCII, otherwise Base64 encoded +// - int64: decimal string representation +// - bool: lowercase "true" or "false" +// +// Values that contain non-ASCII characters, control characters, or +// leading/trailing whitespace are Base64-encoded with the =?base64?...?= wrapper. +// +// The second return value is false if the value is not a supported primitive type. +func encodeHeaderValue(value any) (string, bool) { + s, ok := primitiveToString(value) + if !ok { + return "", false + } + if requiresBase64Encoding(s) { + return encodeBase64(s), true + } + return s, true +} + +// decodeHeaderValue decodes a header value that may be Base64-encoded +// with the =?base64?...?= wrapper. +// +// The second return value is false if the header value is not a valid Base64 encoded value. +func decodeHeaderValue(headerValue string) (string, bool) { + if len(headerValue) == 0 { + return headerValue, true + } + + if encoded, ok := strings.CutPrefix(headerValue, base64Prefix); ok { + if encoded, ok = strings.CutSuffix(encoded, base64Suffix); ok { + decoded, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + return "", false + } + return string(decoded), true + } + } + return headerValue, true +} + +func requiresBase64Encoding(s string) bool { + if len(s) == 0 { + return false + } + if s[0] == ' ' || s[0] == '\t' || s[len(s)-1] == ' ' || s[len(s)-1] == '\t' { + return true + } + for _, c := range s { + if c < 0x20 || c > 0x7E { + return true + } + } + // Per SEP-2243, plain-ASCII values that match the base64 sentinel pattern + // must also be base64-encoded to avoid ambiguity with already-encoded values. + if strings.HasPrefix(s, base64Prefix) && strings.HasSuffix(s, base64Suffix) { + return true + } + return false +} + +func encodeBase64(s string) string { + return base64Prefix + base64.StdEncoding.EncodeToString([]byte(s)) + base64Suffix +} diff --git a/mcp/streamable_headers_test.go b/mcp/streamable_headers_test.go new file mode 100644 index 0000000..97af67c --- /dev/null +++ b/mcp/streamable_headers_test.go @@ -0,0 +1,1649 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package mcp + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "strings" + "testing" + + "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" +) + +func TestExtractName(t *testing.T) { + tests := []struct { + name string + method string + params json.RawMessage + wantName string + wantOK bool + }{ + { + name: "tools/call", + method: "tools/call", + params: mustMarshal(&CallToolParams{Name: "my-tool"}), + wantName: "my-tool", + wantOK: true, + }, + { + name: "prompts/get", + method: "prompts/get", + params: mustMarshal(&GetPromptParams{Name: "code_review"}), + wantName: "code_review", + wantOK: true, + }, + { + name: "resources/read", + method: "resources/read", + params: mustMarshal(&ReadResourceParams{URI: "file:///info.txt"}), + wantName: "file:///info.txt", + wantOK: true, + }, + { + name: "tools/call with empty name", + method: "tools/call", + params: mustMarshal(&CallToolParams{Name: ""}), + wantName: "", + wantOK: true, + }, + { + name: "tool name with hyphen", + method: "tools/call", + params: mustMarshal(&CallToolParams{Name: "my-tool-v2"}), + wantName: "my-tool-v2", + wantOK: true, + }, + { + name: "tool name with underscore", + method: "tools/call", + params: mustMarshal(&CallToolParams{Name: "my_tool_name"}), + wantName: "my_tool_name", + wantOK: true, + }, + { + name: "resource URI with special chars", + method: "resources/read", + params: mustMarshal(&ReadResourceParams{URI: "file:///path/to/file%20name.txt"}), + wantName: "file:///path/to/file%20name.txt", + wantOK: true, + }, + { + name: "resource URI with query string", + method: "resources/read", + params: mustMarshal(&ReadResourceParams{URI: "https://example.com/resource?id=123"}), + wantName: "https://example.com/resource?id=123", + wantOK: true, + }, + { + name: "unrelated method", + method: "initialize", + params: mustMarshal(&InitializeParams{ProtocolVersion: "2025-06-18"}), + wantName: "", + wantOK: false, + }, + { + name: "notification method", + method: "notifications/initialized", + params: nil, + wantName: "", + wantOK: false, + }, + { + name: "invalid JSON params", + method: "tools/call", + params: []byte("not json"), + wantName: "", + wantOK: false, + }, + { + name: "nil params", + method: "tools/call", + params: nil, + wantName: "", + wantOK: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotName, gotOK := extractName(tt.method, tt.params) + if gotName != tt.wantName || gotOK != tt.wantOK { + t.Errorf("extractName(%q, ...) = (%q, %v), want (%q, %v)", + tt.method, gotName, gotOK, tt.wantName, tt.wantOK) + } + }) + } +} + +func TestSetStandardHeaders(t *testing.T) { + tests := []struct { + name string + protocolVersion string + msg jsonrpc.Message + wantMethodHeader string + wantNameHeader string + }{ + { + name: "tools/call with future version", + protocolVersion: minVersionForStandardHeaders, + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantMethodHeader: "tools/call", + wantNameHeader: "my-tool", + }, + { + name: "prompts/get with future version", + protocolVersion: minVersionForStandardHeaders, + msg: &jsonrpc.Request{Method: "prompts/get", Params: mustMarshal(&GetPromptParams{Name: "code_review"})}, + wantMethodHeader: "prompts/get", + wantNameHeader: "code_review", + }, + { + name: "resources/read with future version", + protocolVersion: minVersionForStandardHeaders, + msg: &jsonrpc.Request{Method: "resources/read", Params: mustMarshal(&ReadResourceParams{URI: "file:///info.txt"})}, + wantMethodHeader: "resources/read", + wantNameHeader: "file:///info.txt", + }, + { + name: "initialize sets method only", + protocolVersion: minVersionForStandardHeaders, + msg: &jsonrpc.Request{Method: "initialize", Params: mustMarshal(&InitializeParams{ProtocolVersion: minVersionForStandardHeaders})}, + wantMethodHeader: "initialize", + wantNameHeader: "", + }, + { + name: "notification sets method only", + protocolVersion: minVersionForStandardHeaders, + msg: &jsonrpc.Request{Method: "notifications/initialized"}, + wantMethodHeader: "notifications/initialized", + wantNameHeader: "", + }, + { + name: "old version skips all headers", + protocolVersion: protocolVersion20251125, + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantMethodHeader: "", + wantNameHeader: "", + }, + { + name: "empty version skips all headers", + protocolVersion: "", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantMethodHeader: "", + wantNameHeader: "", + }, + { + name: "nil message is a no-op", + protocolVersion: minVersionForStandardHeaders, + msg: nil, + wantMethodHeader: "", + wantNameHeader: "", + }, + { + name: "response message is ignored", + protocolVersion: minVersionForStandardHeaders, + msg: &jsonrpc.Response{}, + wantMethodHeader: "", + wantNameHeader: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + header := http.Header{} + if tt.protocolVersion != "" { + header.Set(protocolVersionHeader, tt.protocolVersion) + } + + setStandardHeaders(context.Background(), header, tt.msg) + + if got := header.Get(methodHeader); got != tt.wantMethodHeader { + t.Errorf("MethodHeader = %q, want %q", got, tt.wantMethodHeader) + } + if got := header.Get(nameHeader); got != tt.wantNameHeader { + t.Errorf("NameHeader = %q, want %q", got, tt.wantNameHeader) + } + }) + } +} + +func TestValidateMcpHeaders(t *testing.T) { + + tests := []struct { + name string + version string + methodHeader string + nameHeader string + msg jsonrpc.Message + wantErr bool + wantErrContain string + }{ + { + name: "old version skips validation", + version: protocolVersion20251125, + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantErr: false, + }, + { + name: "empty version skips validation", + version: "", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantErr: false, + }, + { + name: "missing Mcp-Method header", + version: minVersionForStandardHeaders, + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantErr: true, + wantErrContain: "missing required Mcp-Method header", + }, + { + name: "missing Mcp-Name for tools/call", + version: minVersionForStandardHeaders, + methodHeader: "tools/call", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantErr: true, + wantErrContain: "missing required Mcp-Name header", + }, + { + name: "missing Mcp-Name for resources/read", + version: minVersionForStandardHeaders, + methodHeader: "resources/read", + msg: &jsonrpc.Request{Method: "resources/read", Params: mustMarshal(&ReadResourceParams{URI: "file:///info.txt"})}, + wantErr: true, + wantErrContain: "missing required Mcp-Name header", + }, + { + name: "missing Mcp-Name for prompts/get", + version: minVersionForStandardHeaders, + methodHeader: "prompts/get", + msg: &jsonrpc.Request{Method: "prompts/get", Params: mustMarshal(&GetPromptParams{Name: "review"})}, + wantErr: true, + wantErrContain: "missing required Mcp-Name header", + }, + { + name: "method mismatch", + version: minVersionForStandardHeaders, + methodHeader: "tools/call", + msg: &jsonrpc.Request{Method: "prompts/get", Params: mustMarshal(&GetPromptParams{Name: "review"})}, + wantErr: true, + wantErrContain: "Mcp-Method header value 'tools/call' does not match body value 'prompts/get'", + }, + { + name: "tool name mismatch", + version: minVersionForStandardHeaders, + methodHeader: "tools/call", + nameHeader: "wrong-tool", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "right-tool"})}, + wantErr: true, + wantErrContain: "Mcp-Name header value 'wrong-tool' does not match body value 'right-tool'", + }, + { + name: "resource URI mismatch", + version: minVersionForStandardHeaders, + methodHeader: "resources/read", + nameHeader: "file:///wrong.txt", + msg: &jsonrpc.Request{Method: "resources/read", Params: mustMarshal(&ReadResourceParams{URI: "file:///right.txt"})}, + wantErr: true, + wantErrContain: "Mcp-Name header value 'file:///wrong.txt' does not match body value 'file:///right.txt'", + }, + { + name: "prompt name mismatch", + version: minVersionForStandardHeaders, + methodHeader: "prompts/get", + nameHeader: "wrong-prompt", + msg: &jsonrpc.Request{Method: "prompts/get", Params: mustMarshal(&GetPromptParams{Name: "right-prompt"})}, + wantErr: true, + wantErrContain: "Mcp-Name header value 'wrong-prompt' does not match body value 'right-prompt'", + }, + { + name: "method value is case-sensitive", + version: minVersionForStandardHeaders, + methodHeader: "TOOLS/CALL", + nameHeader: "my-tool", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantErr: true, + wantErrContain: "Mcp-Method header value 'TOOLS/CALL' does not match body value 'tools/call'", + }, + { + name: "valid tools/call", + version: minVersionForStandardHeaders, + methodHeader: "tools/call", + nameHeader: "my-tool", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantErr: false, + }, + { + name: "valid resources/read", + version: minVersionForStandardHeaders, + methodHeader: "resources/read", + nameHeader: "file:///info.txt", + msg: &jsonrpc.Request{Method: "resources/read", Params: mustMarshal(&ReadResourceParams{URI: "file:///info.txt"})}, + wantErr: false, + }, + { + name: "valid prompts/get", + version: minVersionForStandardHeaders, + methodHeader: "prompts/get", + nameHeader: "code_review", + msg: &jsonrpc.Request{Method: "prompts/get", Params: mustMarshal(&GetPromptParams{Name: "code_review"})}, + wantErr: false, + }, + { + name: "Base64-encoded Mcp-Name is decoded before comparison", + version: minVersionForStandardHeaders, + methodHeader: "tools/call", + nameHeader: base64Prefix + base64.StdEncoding.EncodeToString([]byte("café-tool")) + base64Suffix, + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "café-tool"})}, + wantErr: false, + }, + { + name: "Mcp-Name with invalid Base64 encoding", + version: minVersionForStandardHeaders, + methodHeader: "tools/call", + nameHeader: base64Prefix + "not valid base64!" + base64Suffix, + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantErr: true, + wantErrContain: "invalid Base64 encoding", + }, + { + name: "valid initialize (no name needed)", + version: minVersionForStandardHeaders, + methodHeader: "initialize", + msg: &jsonrpc.Request{Method: "initialize", Params: mustMarshal(&InitializeParams{ProtocolVersion: minVersionForStandardHeaders})}, + wantErr: false, + }, + { + name: "valid notification (no name needed)", + version: minVersionForStandardHeaders, + methodHeader: "notifications/initialized", + msg: &jsonrpc.Request{Method: "notifications/initialized"}, + wantErr: false, + }, + { + name: "tool name with hyphen", + version: minVersionForStandardHeaders, + methodHeader: "tools/call", + nameHeader: "my-tool-name", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool-name"})}, + wantErr: false, + }, + { + name: "tool name with underscore", + version: minVersionForStandardHeaders, + methodHeader: "tools/call", + nameHeader: "my_tool_name", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my_tool_name"})}, + wantErr: false, + }, + { + name: "resource URI with special chars", + version: minVersionForStandardHeaders, + methodHeader: "resources/read", + nameHeader: "file:///path/to/file%20name.txt", + msg: &jsonrpc.Request{Method: "resources/read", Params: mustMarshal(&ReadResourceParams{URI: "file:///path/to/file%20name.txt"})}, + wantErr: false, + }, + { + name: "resource URI with query string", + version: minVersionForStandardHeaders, + methodHeader: "resources/read", + nameHeader: "https://example.com/resource?id=123", + msg: &jsonrpc.Request{Method: "resources/read", Params: mustMarshal(&ReadResourceParams{URI: "https://example.com/resource?id=123"})}, + wantErr: false, + }, + { + name: "response message is ignored", + version: minVersionForStandardHeaders, + msg: &jsonrpc.Response{}, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + header := http.Header{} + if tt.version != "" { + header.Set(protocolVersionHeader, tt.version) + } + if tt.methodHeader != "" { + header.Set(methodHeader, tt.methodHeader) + } + if tt.nameHeader != "" { + header.Set(nameHeader, tt.nameHeader) + } + + err := validateMcpHeaders(header, tt.msg, nil) + if tt.wantErr { + if err == nil { + t.Fatalf("validateMcpHeaders() = nil, want error containing %q", tt.wantErrContain) + } + if !strings.Contains(err.Error(), tt.wantErrContain) { + t.Errorf("validateMcpHeaders() error = %q, want substring %q", err.Error(), tt.wantErrContain) + } + } else if err != nil { + t.Errorf("validateMcpHeaders() = %v, want nil", err) + } + }) + } +} + +func TestValidateToolParamHeaders(t *testing.T) { + tests := []struct { + name string + tool *Tool + wantErr bool + wantErrSub string + }{ + { + name: "valid tool with x-mcp-header", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "Region", + }, + }, + }, + }, + }, + { + name: "tool with no x-mcp-header annotations", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "query": map[string]any{"type": "string"}, + }, + }, + }, + }, + { + name: "empty x-mcp-header value", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "", + }, + }, + }, + }, + wantErr: true, + wantErrSub: "non-empty string", + }, + { + name: "x-mcp-header with space", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "My Region", + }, + }, + }, + }, + wantErr: true, + wantErrSub: "invalid character", + }, + { + name: "x-mcp-header with colon", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "Region:Primary", + }, + }, + }, + }, + wantErr: true, + wantErrSub: "invalid character", + }, + { + name: "x-mcp-header with non-ASCII", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "Région", + }, + }, + }, + }, + wantErr: true, + wantErrSub: "invalid character", + }, + { + name: "x-mcp-header with separator char (parens) is invalid per RFC 9110", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "X-(Region)", + }, + }, + }, + }, + wantErr: true, + wantErrSub: "invalid character", + }, + { + name: "x-mcp-header with equals sign is invalid per RFC 9110", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "Region=1", + }, + }, + }, + }, + wantErr: true, + wantErrSub: "invalid character", + }, + { + name: "x-mcp-header with all tchar specials is valid", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "!#$%&'*+-.^_`|~aZ0", + }, + }, + }, + }, + }, + { + name: "duplicate header names same case", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "a": map[string]any{"type": "string", "x-mcp-header": "Region"}, + "b": map[string]any{"type": "string", "x-mcp-header": "Region"}, + }, + }, + }, + wantErr: true, + wantErrSub: "duplicate", + }, + { + name: "duplicate header names different case", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "a": map[string]any{"type": "string", "x-mcp-header": "Region"}, + "b": map[string]any{"type": "string", "x-mcp-header": "REGION"}, + }, + }, + }, + wantErr: true, + wantErrSub: "duplicate", + }, + { + name: "x-mcp-header on array type", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "items": map[string]any{ + "type": "array", + "x-mcp-header": "Items", + }, + }, + }, + }, + wantErr: true, + wantErrSub: "primitive types", + }, + { + name: "x-mcp-header on object type", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "nested": map[string]any{ + "type": "object", + "x-mcp-header": "Nested", + }, + }, + }, + }, + wantErr: true, + wantErrSub: "primitive types", + }, + { + name: "x-mcp-header on number type is invalid", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "count": map[string]any{ + "type": "number", + "x-mcp-header": "Count", + }, + }, + }, + }, + wantErr: true, + wantErrSub: "primitive types", + }, + { + name: "x-mcp-header on integer type is valid", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "count": map[string]any{ + "type": "integer", + "x-mcp-header": "Count", + }, + }, + }, + }, + }, + { + name: "x-mcp-header on boolean type is valid", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "flag": map[string]any{ + "type": "boolean", + "x-mcp-header": "Flag", + }, + }, + }, + }, + }, + { + name: "x-mcp-header on nested property inside object is valid", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config": map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "Region", + }, + }, + }, + }, + }, + }, + }, + { + name: "x-mcp-header on deeply nested property is valid", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "outer": map[string]any{ + "type": "object", + "properties": map[string]any{ + "inner": map[string]any{ + "type": "object", + "properties": map[string]any{ + "value": map[string]any{ + "type": "string", + "x-mcp-header": "Value", + }, + }, + }, + }, + }, + }, + }, + }, + }, + { + name: "duplicate header names across nesting levels", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{"type": "string", "x-mcp-header": "Region"}, + "config": map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{"type": "string", "x-mcp-header": "Region"}, + }, + }, + }, + }, + }, + wantErr: true, + wantErrSub: "duplicate", + }, + { + name: "object property without nested x-mcp-header is valid", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config": map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + }, + }, + }, + "flag": map[string]any{ + "type": "boolean", + "x-mcp-header": "Flag", + }, + }, + }, + }, + }, + { + name: "jsonschema.Schema valid x-mcp-header", + tool: &Tool{ + Name: "test", + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "region": { + Type: "string", + Extra: map[string]any{"x-mcp-header": "Region"}, + }, + }, + }, + }, + }, + { + name: "jsonschema.Schema x-mcp-header on array type", + tool: &Tool{ + Name: "test", + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "items": { + Type: "array", + Extra: map[string]any{"x-mcp-header": "Items"}, + }, + }, + }, + }, + wantErr: true, + wantErrSub: "primitive types", + }, + { + name: "jsonschema.Schema nested x-mcp-header is valid", + tool: &Tool{ + Name: "test", + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "config": { + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "region": { + Type: "string", + Extra: map[string]any{"x-mcp-header": "Region"}, + }, + }, + }, + }, + }, + }, + }, + { + name: "json.RawMessage valid x-mcp-header", + tool: &Tool{ + Name: "test", + InputSchema: json.RawMessage(`{"type":"object","properties":{"region":{"type":"string","x-mcp-header":"Region"}}}`), + }, + }, + { + name: "json.RawMessage x-mcp-header on object type", + tool: &Tool{ + Name: "test", + InputSchema: json.RawMessage(`{"type":"object","properties":{"nested":{"type":"object","x-mcp-header":"Nested"}}}`), + }, + wantErr: true, + wantErrSub: "primitive types", + }, + { + name: "nested invalid header name is rejected", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config": map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "Bad Header", // contains a space + }, + }, + }, + }, + }, + }, + wantErr: true, + wantErrSub: "invalid character", + }, + { + name: "nested x-mcp-header on number type is rejected", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config": map[string]any{ + "type": "object", + "properties": map[string]any{ + "count": map[string]any{ + "type": "number", + "x-mcp-header": "Count", + }, + }, + }, + }, + }, + }, + wantErr: true, + wantErrSub: "primitive types", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateParamHeaderAnnotations(tt.tool) + if tt.wantErr { + if err == nil { + t.Fatal("validateToolParamHeaders() = nil, want error") + } + if !strings.Contains(err.Error(), tt.wantErrSub) { + t.Errorf("error = %q, want substring %q", err.Error(), tt.wantErrSub) + } + } else if err != nil { + t.Errorf("validateToolParamHeaders() = %v, want nil", err) + } + }) + } +} + +func TestFilterValidTools(t *testing.T) { + valid := &Tool{ + Name: "valid", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{"type": "string", "x-mcp-header": "Region"}, + }, + }, + } + invalid := &Tool{ + Name: "invalid", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{"type": "string", "x-mcp-header": ""}, + }, + }, + } + noAnnotation := &Tool{ + Name: "plain", + InputSchema: map[string]any{"type": "object", "properties": map[string]any{"q": map[string]any{"type": "string"}}}, + } + nestedValid := &Tool{ + Name: "nested-valid", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config": map[string]any{ + "type": "object", + "properties": map[string]any{ + "tenant": map[string]any{"type": "string", "x-mcp-header": "TenantId"}, + }, + }, + }, + }, + } + + validJsonSchema := &Tool{ + Name: "valid-jsonschema", + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "region": { + Type: "string", + Extra: map[string]any{"x-mcp-header": "Region"}, + }, + }, + }, + } + invalidJsonSchema := &Tool{ + Name: "invalid-jsonschema", + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "items": { + Type: "array", + Extra: map[string]any{"x-mcp-header": "Items"}, + }, + }, + }, + } + + result := filterValidTools(nil, []*Tool{valid, invalid, noAnnotation, nestedValid, validJsonSchema, invalidJsonSchema}) + if len(result) != 4 { + t.Fatalf("filterValidTools returned %d tools, want 4", len(result)) + } + if result[0].Name != "valid" || result[1].Name != "plain" || result[2].Name != "nested-valid" || result[3].Name != "valid-jsonschema" { + t.Errorf("filterValidTools returned [%s, %s, %s, %s], want [valid, plain, nested-valid, valid-jsonschema]", + result[0].Name, result[1].Name, result[2].Name, result[3].Name) + } + + // Regression for #1119: a malformed "tools":[null] response must not panic; + // nil tools are treated as invalid and excluded. + if got := filterValidTools(nil, []*Tool{nil}); len(got) != 0 { + t.Errorf("filterValidTools([nil]) returned %d tools, want 0", len(got)) + } + if got := filterValidTools(nil, []*Tool{nil, valid, nil}); len(got) != 1 || got[0].Name != "valid" { + t.Errorf("filterValidTools([nil, valid, nil]) = %d tools, want [valid]", len(got)) + } +} + +func TestSetStandardHeadersWithParamHeaders(t *testing.T) { + toolSchema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "Region", + }, + "query": map[string]any{ + "type": "string", + }, + "priority": map[string]any{ + "type": "string", + "x-mcp-header": "Priority", + }, + }, + } + tool := &Tool{Name: "execute_sql", InputSchema: toolSchema} + + tests := []struct { + name string + tool *Tool + params any + wantHeaders map[string]string + }{ + { + name: "sets param headers from arguments", + tool: tool, + params: &CallToolParams{ + Name: "execute_sql", + Arguments: map[string]any{"region": "us-west1", "query": "SELECT 1", "priority": "high"}, + }, + wantHeaders: map[string]string{ + "Mcp-Param-Region": "us-west1", + "Mcp-Param-Priority": "high", + }, + }, + { + name: "omits header when argument is missing", + tool: tool, + params: &CallToolParams{ + Name: "execute_sql", + Arguments: map[string]any{"query": "SELECT 1"}, + }, + wantHeaders: nil, + }, + { + name: "omits header when argument is null", + tool: tool, + params: &CallToolParams{ + Name: "execute_sql", + Arguments: map[string]any{"region": nil, "query": "SELECT 1"}, + }, + wantHeaders: nil, + }, + { + name: "encodes non-ASCII value", + tool: tool, + params: &CallToolParams{ + Name: "execute_sql", + Arguments: map[string]any{"region": "日本", "query": "SELECT 1"}, + }, + wantHeaders: map[string]string{ + "Mcp-Param-Region": "=?base64?5pel5pys?=", + }, + }, + { + name: "handles boolean argument", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "flag": map[string]any{"type": "boolean", "x-mcp-header": "Flag"}, + }, + }, + }, + params: &CallToolParams{ + Name: "test", + Arguments: map[string]any{"flag": true}, + }, + wantHeaders: map[string]string{ + "Mcp-Param-Flag": "true", + }, + }, + { + name: "handles integer argument", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "count": map[string]any{"type": "integer", "x-mcp-header": "Count"}, + }, + }, + }, + params: &CallToolParams{ + Name: "test", + Arguments: map[string]any{"count": float64(42)}, + }, + wantHeaders: map[string]string{ + "Mcp-Param-Count": "42", + }, + }, + { + name: "out-of-range integer argument produces no header", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "count": map[string]any{"type": "integer", "x-mcp-header": "Count"}, + }, + }, + }, + params: &CallToolParams{ + Name: "test", + Arguments: map[string]any{"count": float64(maxSafeInteger) + 2}, + }, + wantHeaders: nil, + }, + { + name: "no tool in extra does not add param headers", + tool: nil, + params: &CallToolParams{ + Name: "execute_sql", + Arguments: map[string]any{"region": "us-west1"}, + }, + wantHeaders: nil, + }, + { + name: "nested arguments resolve via dotted path", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config": map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{"type": "string", "x-mcp-header": "Region"}, + "tenant": map[string]any{"type": "string", "x-mcp-header": "TenantId"}, + }, + }, + }, + }, + }, + params: &CallToolParams{ + Name: "test", + Arguments: map[string]any{ + "config": map[string]any{ + "region": "us-west1", + "tenant": "acme", + }, + }, + }, + wantHeaders: map[string]string{ + "Mcp-Param-Region": "us-west1", + "Mcp-Param-TenantId": "acme", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + header := http.Header{} + header.Set(protocolVersionHeader, minVersionForStandardHeaders) + + ctx := context.Background() + if tt.tool != nil { + ctx = context.WithValue(ctx, toolContextKey, tt.tool) + } + + msg := &jsonrpc.Request{ + Method: "tools/call", + Params: mustMarshal(tt.params), + } + + setStandardHeaders(ctx, header, msg) + + if got := header.Get(methodHeader); got != "tools/call" { + t.Errorf("MethodHeader = %q, want %q", got, "tools/call") + } + + for h, want := range tt.wantHeaders { + if got := header.Get(h); got != want { + t.Errorf("%s = %q, want %q", h, got, want) + } + } + + if got := header.Get("Mcp-Param-query"); got != "" { + t.Errorf("non-annotated param got header: Mcp-Param-query = %q", got) + } + }) + } +} + +func TestExtractToolParamHeaders(t *testing.T) { + tests := []struct { + name string + tool *Tool + want map[string]string + }{ + { + name: "extracts x-mcp-header annotations", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{"type": "string", "x-mcp-header": "Region"}, + "query": map[string]any{"type": "string"}, + "tenant_id": map[string]any{"type": "string", "x-mcp-header": "TenantId"}, + }, + }, + }, + want: map[string]string{"region": "Region", "tenant_id": "TenantId"}, + }, + { + name: "returns nil for tool without properties", + tool: &Tool{Name: "test", InputSchema: map[string]any{"type": "object"}}, + want: nil, + }, + { + name: "returns nil for non-map schema", + tool: &Tool{Name: "test", InputSchema: "not a map"}, + want: nil, + }, + { + name: "returns nil when no annotations", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{"q": map[string]any{"type": "string"}}, + }, + }, + want: nil, + }, + { + name: "jsonschema.Schema with x-mcp-header in Extra", + tool: &Tool{ + Name: "test", + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "region": { + Type: "string", + Extra: map[string]any{"x-mcp-header": "Region"}, + }, + "query": {Type: "string"}, + }, + }, + }, + want: map[string]string{"region": "Region"}, + }, + { + name: "json.RawMessage with x-mcp-header", + tool: &Tool{ + Name: "test", + InputSchema: json.RawMessage(`{"type":"object","properties":{"region":{"type":"string","x-mcp-header":"Region"},"query":{"type":"string"}}}`), + }, + want: map[string]string{"region": "Region"}, + }, + { + name: "jsonschema.Schema without x-mcp-header", + tool: &Tool{ + Name: "test", + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "query": {Type: "string"}, + }, + }, + }, + want: nil, + }, + { + name: "nested x-mcp-header annotations produce path-slice bindings", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{"type": "string", "x-mcp-header": "Region"}, + "config": map[string]any{ + "type": "object", + "properties": map[string]any{ + "tenant": map[string]any{"type": "string", "x-mcp-header": "TenantId"}, + "deep": map[string]any{ + "type": "object", + "properties": map[string]any{ + "flag": map[string]any{"type": "boolean", "x-mcp-header": "DeepFlag"}, + }, + }, + }, + }, + }, + }, + }, + want: map[string]string{ + "region": "Region", + "config.tenant": "TenantId", + "config.deep.flag": "DeepFlag", + }, + }, + { + name: "property name containing a dot is preserved in path", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "user.id": map[string]any{"type": "string", "x-mcp-header": "UserId"}, + }, + }, + }, + want: map[string]string{"user.id": "UserId"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := extractParamHeaderAnnotations(tt.tool) + if tt.want == nil { + if got != nil { + t.Errorf("extractToolParamHeaders() = %v, want nil", got) + } + return + } + if len(got) != len(tt.want) { + t.Fatalf("extractToolParamHeaders() returned %d entries, want %d", len(got), len(tt.want)) + } + // Index returned bindings by joined-path for comparison; the + // expected map uses dotted-path keys for readability. + gotMap := make(map[string]string, len(got)) + for _, b := range got { + gotMap[strings.Join(b.Path, ".")] = b.Header + } + for k, v := range tt.want { + if gotMap[k] != v { + t.Errorf("extractToolParamHeaders()[%q] = %q, want %q", k, gotMap[k], v) + } + } + }) + } +} + +func TestUnmarshalPrimitive(t *testing.T) { + tests := []struct { + name string + raw string + want any + }{ + {"string", `"hello"`, "hello"}, + {"true", `true`, true}, + {"false", `false`, false}, + + // Integer JSON numbers are promoted to int64. + {"integer", `42`, int64(42)}, + {"integer zero", `0`, int64(0)}, + {"integer negative", `-7`, int64(-7)}, + {"integer max safe", `9007199254740991`, int64(maxSafeInteger)}, + {"integer min safe", `-9007199254740991`, int64(minSafeInteger)}, + // JSON serialization of an integer-valued float (e.g. "42.0") is + // still a valid integer at the value level and must be accepted. + {"integer-valued float", `42.0`, int64(42)}, + + // Non-integer numbers, out-of-range integers, and disallowed JSON + // kinds are rejected (return nil). + {"float with fraction", `3.14`, nil}, + {"integer above max safe", `9007199254740993`, nil}, + {"integer below min safe", `-9007199254740993`, nil}, + {"null", `null`, nil}, + {"array", `[1,2]`, nil}, + {"object", `{"a":1}`, nil}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := unmarshalPrimitive(json.RawMessage(tt.raw)) + if fmt.Sprintf("%v", got) != fmt.Sprintf("%v", tt.want) { + t.Errorf("unmarshalPrimitive(%s) = %v (%T), want %v (%T)", tt.raw, got, got, tt.want, tt.want) + } + }) + } +} + +func TestEncodeHeaderValue(t *testing.T) { + tests := []struct { + name string + value any + want string + wantOK bool + }{ + // Strings + {"plain ASCII", "us-west1", "us-west1", true}, + {"empty string", "", "", true}, + {"string with internal spaces", "us west 1", "us west 1", true}, + {"string with leading space", " us-west1", "=?base64?IHVzLXdlc3Qx?=", true}, + {"string with trailing space", "us-west1 ", "=?base64?dXMtd2VzdDEg?=", true}, + {"string with both spaces", " us-west1 ", "=?base64?IHVzLXdlc3QxIA==?=", true}, + {"non-ASCII", "日本語", "=?base64?5pel5pys6Kqe?=", true}, + {"mixed ASCII and non-ASCII", "Hello, 世界", "=?base64?SGVsbG8sIOS4lueVjA==?=", true}, + {"string with newline", "line1\nline2", "=?base64?bGluZTEKbGluZTI=?=", true}, + {"string with carriage return", "line1\r\nline2", "=?base64?bGluZTENCmxpbmUy?=", true}, + {"string with leading tab", "\tindented", "=?base64?CWluZGVudGVk?=", true}, + + // Sentinel pattern collisions: plain-ASCII values that match the base64 + // sentinel pattern must also be base64-encoded to avoid ambiguity. + {"sentinel collision literal", "=?base64?literal?=", "=?base64?PT9iYXNlNjQ/bGl0ZXJhbD89?=", true}, + {"sentinel collision empty", "=?base64??=", "=?base64?PT9iYXNlNjQ/Pz0=?=", true}, + // Uppercase sentinel does NOT collide (case-sensitive markers). + {"uppercase pseudo-sentinel passes through", "=?BASE64?abc?=", "=?BASE64?abc?=", true}, + + {"integer", int64(42), "42", true}, + {"integer zero", int64(0), "0", true}, + {"integer negative", int64(-7), "-7", true}, + {"integer max safe", int64(maxSafeInteger), "9007199254740991", true}, + {"integer min safe", int64(minSafeInteger), "-9007199254740991", true}, + + // Booleans + {"true", true, "true", true}, + {"false", false, "false", true}, + + {"raw float64 rejected", float64(42), "", false}, + {"nil", nil, "", false}, + {"slice", []string{"a"}, "", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := encodeHeaderValue(tt.value) + if ok != tt.wantOK { + t.Fatalf("encodeHeaderValue(%v) ok = %v, want %v", tt.value, ok, tt.wantOK) + } + if got != tt.want { + t.Errorf("encodeHeaderValue(%v) = %q, want %q", tt.value, got, tt.want) + } + }) + } +} + +func TestDecodeHeaderValue(t *testing.T) { + tests := []struct { + name string + input string + want string + wantOK bool + }{ + {"plain value", "us-west1", "us-west1", true}, + {"empty value", "", "", true}, + {"valid base64", "=?base64?SGVsbG8=?=", "Hello", true}, + {"non-ASCII decoded", "=?base64?5pel5pys6Kqe?=", "日本語", true}, + {"leading space decoded", "=?base64?IHVzLXdlc3Qx?=", " us-west1", true}, + // Per SEP-2243, the base64 sentinel markers are case-sensitive: an + // uppercase prefix is treated as a literal value, not a base64 marker. + {"uppercase prefix is literal", "=?BASE64?SGVsbG8=?=", "=?BASE64?SGVsbG8=?=", true}, + {"mixed case prefix is literal", "=?Base64?SGVsbG8=?=", "=?Base64?SGVsbG8=?=", true}, + {"invalid base64 chars", "=?base64?SGVs!!!bG8=?=", "", false}, + // Missing prefix or suffix: treated as literal values, not base64 + {"missing prefix", "SGVsbG8=", "SGVsbG8=", true}, + {"missing suffix", "=?base64?SGVsbG8=", "=?base64?SGVsbG8=", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := decodeHeaderValue(tt.input) + if ok != tt.wantOK { + t.Fatalf("decodeHeaderValue(%q) ok = %v, want %v", tt.input, ok, tt.wantOK) + } + if got != tt.want { + t.Errorf("decodeHeaderValue(%q) = %q, want %q", tt.input, got, tt.want) + } + }) + } +} + +func TestEncodeDecodeRoundTrip(t *testing.T) { + values := []string{ + "us-west1", + "", + " leading", + "trailing ", + "Hello, 世界", + "line1\nline2", + "\ttab", + "=?base64?literal?=", // sentinel-pattern collision (SEP-2243) + } + for _, v := range values { + encoded, ok := encodeHeaderValue(v) + if !ok { + t.Fatalf("encodeHeaderValue(%q) failed", v) + } + decoded, ok := decodeHeaderValue(encoded) + if !ok { + t.Fatalf("decodeHeaderValue(%q) failed", encoded) + } + if decoded != v { + t.Errorf("round-trip failed: %q -> %q -> %q", v, encoded, decoded) + } + } +} + +// TestValidateParamHeaders_IntegerComparison verifies server-side validation +// of integer x-mcp-header parameters per SEP-2243. +func TestValidateParamHeaders_IntegerComparison(t *testing.T) { + tool := &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "count": map[string]any{"type": "integer", "x-mcp-header": "Count"}, + }, + }, + } + tests := []struct { + name string + headerVal string + bodyArg any + wantErr bool + }{ + // Canonical decimal form matches. + {"integer matches integer", "42", float64(42), false}, + {"integer header matches integer-valued float body", "42", float64(42.0), false}, + {"negative integer matches", "-7", float64(-7), false}, + {"large safe integer matches", "1000000000000", float64(1e12), false}, + + {"non-canonical '42.0' header matches integer body", "42.0", float64(42), false}, + {"scientific notation header matches integer body", "1e2", float64(100), false}, + {"negative non-canonical header matches", "-7.0", float64(-7), false}, + + // Genuine mismatches and invalid forms are still rejected. + {"different integers do not match", "42", float64(43), true}, + {"fractional header against integer body", "3.14", float64(3), true}, + {"non-numeric header fails", "not-a-number", float64(42), true}, + {"header outside safe range against integer body", "9007199254740993", float64(42), true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + header := http.Header{} + header.Set(paramHeaderPrefix+"Count", tt.headerVal) + args := map[string]any{"count": tt.bodyArg} + msg := &jsonrpc.Request{ + Method: "tools/call", + Params: mustMarshal(&CallToolParams{Name: "test", Arguments: args}), + } + err := validateParamHeaders(header, msg, tool) + if tt.wantErr && err == nil { + t.Errorf("validateParamHeaders() = nil, want error") + } + if !tt.wantErr && err != nil { + t.Errorf("validateParamHeaders() = %v, want nil", err) + } + }) + } +} + +// TestValidateParamHeaders_NestedArguments verifies that the server-side +// validation can look up nested arguments via the dotted path produced by +// extractParamHeaderAnnotations. +func TestValidateParamHeaders_NestedArguments(t *testing.T) { + tool := &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config": map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{"type": "string", "x-mcp-header": "Region"}, + }, + }, + }, + }, + } + header := http.Header{} + header.Set(paramHeaderPrefix+"Region", "us-west1") + args := map[string]any{ + "config": map[string]any{"region": "us-west1"}, + } + msg := &jsonrpc.Request{ + Method: "tools/call", + Params: mustMarshal(&CallToolParams{Name: "test", Arguments: args}), + } + if err := validateParamHeaders(header, msg, tool); err != nil { + t.Errorf("validateParamHeaders() = %v, want nil", err) + } + + // Mismatched nested value should fail. + args["config"].(map[string]any)["region"] = "eu-west1" + msg.Params = mustMarshal(&CallToolParams{Name: "test", Arguments: args}) + if err := validateParamHeaders(header, msg, tool); err == nil { + t.Error("validateParamHeaders() = nil, want error for mismatched nested value") + } +} + +// TestMcpNameHeaderEncodingRoundTrip verifies that a name which is not +// header-safe is Base64-wrapped by the client per the 2026-07-28 transport +// spec, and that the server's validation accepts what the client produced. +func TestMcpNameHeaderEncodingRoundTrip(t *testing.T) { + tests := []struct { + name string + msg *jsonrpc.Request + wantHeader string + }{ + { + name: "ascii-safe tool name is sent as-is", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "my-tool"})}, + wantHeader: "my-tool", + }, + { + name: "non-ascii tool name is base64 encoded", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "café-tool"})}, + wantHeader: encodeBase64("café-tool"), + }, + { + name: "prompt name with a leading space is base64 encoded", + msg: &jsonrpc.Request{Method: "prompts/get", Params: mustMarshal(&GetPromptParams{Name: " leading-space"})}, + wantHeader: encodeBase64(" leading-space"), + }, + { + name: "resource uri with a control character is base64 encoded", + msg: &jsonrpc.Request{Method: "resources/read", Params: mustMarshal(&ReadResourceParams{URI: "file:///a\tb.txt"})}, + wantHeader: encodeBase64("file:///a\tb.txt"), + }, + { + name: "name that looks like the sentinel is base64 encoded", + msg: &jsonrpc.Request{Method: "tools/call", Params: mustMarshal(&CallToolParams{Name: "=?base64?abc?="})}, + wantHeader: encodeBase64("=?base64?abc?="), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + header := http.Header{} + header.Set(protocolVersionHeader, minVersionForStandardHeaders) + setStandardHeaders(context.Background(), header, tt.msg) + + if got := header.Get(nameHeader); got != tt.wantHeader { + t.Errorf("Mcp-Name = %q, want %q", got, tt.wantHeader) + } + + // The server must accept the headers the client just generated. + if err := validateMcpHeaders(header, tt.msg, nil); err != nil { + t.Errorf("validateMcpHeaders() on client-generated headers = %v, want nil", err) + } + }) + } +} diff --git a/mcp/streamable_server.go b/mcp/streamable_server.go new file mode 100644 index 0000000..8a573e5 --- /dev/null +++ b/mcp/streamable_server.go @@ -0,0 +1,160 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +// TODO: move server-side streamable HTTP logic from streamable.go to this file. + +package mcp + +/* +Streamable HTTP Server Design + +This document describes the server-side implementation of the MCP streamable +HTTP transport, as defined by the MCP spec: +https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#streamable-http + +# Overview + +The streamable HTTP transport enables MCP communication over HTTP, with +server-sent events (SSE) for server-to-client messages. The implementation +consists of several layered components: + + ┌─────────────────────────────────────────────────────────────────┐ + │ [StreamableHTTPHandler] │ + │ http.Handler that manages sessions and routes HTTP requests │ + └─────────────────────────────────────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────────────────────────────────────┐ + │ [StreamableServerTransport] │ + │ transport implementation, one per session; exposes ServeHTTP │ + └─────────────────────────────────────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────────────────────────────────────┐ + │ [streamableServerConn] │ + │ Connection implementation, handles message routing │ + └─────────────────────────────────────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────────────────────────────────────┐ + │ [stream] │ + │ Logical message channel within a session, may be resumed │ + └─────────────────────────────────────────────────────────────────┘ + +# Sessions + +As with other transports, a session represents a logical MCP connection between +a client and server. In the streamable transport, sessions are identified by a +unique session ID (Mcp-Session-Id header) and persist across multiple HTTP +requests. + +[StreamableHTTPHandler] maintains a map of active sessions ([sessionInfo]), +each containing: + - The [ServerSession] (MCP-level session state) + - The [StreamableServerTransport] (for message I/O) + - Optional timeout management for idle session cleanup + +Sessions are created on the first POST request (typically containing the +initialize request) and destroyed either by: + - Client sending a DELETE request + - Session timeout due to inactivity + - Server explicitly closing the session + +# Streams + +Within a session, there can be multiple concurrent "streams" - logical channels +for message delivery. This is distinct from HTTP streams; a single [stream] may +span multiple HTTP request/response cycles (via resumption). + +There are two types of streams: + +1. Optional standalone SSE stream (id = ""): + - Created when client sends a GET request to the endpoint + - Used for server-initiated messages (requests/notifications to client) + - Persists for the lifetime of the session + - Only one standalone stream per session + +2. Request streams (id = random string): + - Created for each POST request containing JSON-RPC calls + - Used to route responses back to the originating HTTP request + - Completed when all responses have been sent + - Can be resumed via GET with Last-Event-ID if interrupted + +# Message Routing + +When the server writes a message, it must be routed to the correct [stream]: + + - Responses: Routed to the stream that originated the request + - Requests/Notifications made during request handling: Routed to the same + stream as the triggering request (via context) + - Requests/Notifications made outside request handling: Routed to the + standalone SSE stream + +This routing is implemented using: + - [streamableServerConn.requestStreams] maps request IDs to stream IDs + - [idContextKey] is used to store the originating request ID in Context + - [streamableServerConn.streams] maps stream IDs to [stream] objects + +# Stream Resumption + +If an HTTP connection is interrupted (network issues, etc.), clients can +resume a stream by sending a GET request with the Last-Event-ID header. +This requires an [EventStore] to be configured on the server. + + - [EventStore.Open] is called when a new stream is created + - [EventStore.Append] is called for each message written to the stream + - [EventStore.After] is called to replay messages after a given index + - [EventStore.SessionClosed] is called when the session ends + +Event IDs are formatted as "_" to identify both the +stream and position within that stream (see [formatEventID] and [parseEventID]). + +# Stateless Mode + +For simpler deployments, the handler supports "stateless" mode +([StreamableHTTPOptions.Stateless]) where: + - No session ID validation is performed + - Each request creates a temporary session that's closed after the request + - Server-to-client requests are not supported (no way to receive response) + +This mode is useful for simple tool servers that don't need bidirectional +communication. + +# Response Formats + +The server can respond to POST requests in two formats: + +1. text/event-stream (default): Messages sent as SSE events, supports + streaming multiple messages and server-initiated communication during + request handling. + +2. application/json ([StreamableHTTPOptions.JSONResponse]): Single JSON + response, simpler but doesn't support streaming. Server-initiated messages + during request handling go to the standalone SSE stream instead. + +# HTTP Methods + + - POST: Send JSON-RPC messages (requests, responses, notifications) + - GET: Open standalone SSE stream or resume an interrupted stream + - DELETE: Terminate the session + +# Key Implementation Details + +The [stream] struct manages delivery of messages to HTTP responses. + +Fields: + - [stream.w] is the ResponseWriter for the current HTTP response (non-nil indicates claimed) + - [stream.done] is closed to release the hanging HTTP request + - [stream.requests] tracks pending request IDs (stream completes when empty) + +Methods: + - [stream.deliverLocked] delivers a message to the stream + - [stream.close] sends a close event and releases the stream + - [stream.release] releases the stream from the HTTP request, allowing resumption + +[streamableServerConn] handles the [Connection] interface: + - [streamableServerConn.Read] receives messages from the incoming channel (fed by POST handlers) + - [streamableServerConn.Write] routes messages to appropriate streams + - [streamableServerConn.Close] terminates the session and notifies the [EventStore] +*/ diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 0f38a0f..7af9dff 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -2,6 +2,8 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. +//lint:file-ignore SA1019 tests exercise deprecated SEP-2577 APIs (roots, sampling, logging). + package mcp import ( @@ -108,9 +110,7 @@ func TestStreamableTransports(t *testing.T) { } // Now, spin off a goroutine that runs after the sampling request, to // check behavior when the client request has completed. - sampleWG.Add(1) - go func() { - defer sampleWG.Done() + sampleWG.Go(func() { <-sampleDone // Test that sampling requests in the tool context fail outside of // tool handling, but succeed on the background context. @@ -120,7 +120,7 @@ func TestStreamableTransports(t *testing.T) { } { testSample(test) } - }() + }) return &CallToolResult{}, nil, nil }) @@ -182,7 +182,11 @@ func TestStreamableTransports(t *testing.T) { if sid == "" { t.Fatalf("empty session ID") } - if g, w := session.mcpConn.(*streamableClientConn).initializedResult.ProtocolVersion, latestProtocolVersion; g != w { + // Stateful HTTP servers reject 2026-07-28 (which requires + // StreamableHTTPOptions.Stateless = true), so the client falls + // back to the legacy initialize handshake which explicitly + // requests 2025-11-25. + if g, w := session.mcpConn.(*streamableClientConn).initializedResult.ProtocolVersion, protocolVersion20251125; g != w { t.Fatalf("got protocol version %q, want %q", g, w) } @@ -200,7 +204,10 @@ func TestStreamableTransports(t *testing.T) { if g := session.ID(); g != sid { t.Errorf("session ID: got %q, want %q", g, sid) } - if g, w := lastHeader.Get(protocolVersionHeader), latestProtocolVersion; g != w { + // Stateful HTTP requires the legacy protocol (see the earlier + // initializeResult assertion); the client sends the matching + // MCP-Protocol-Version header. + if g, w := lastHeader.Get(protocolVersionHeader), protocolVersion20251125; g != w { t.Errorf("got protocol version header %q, want %q", g, w) } want := &CallToolResult{ @@ -268,9 +275,7 @@ func TestStreamableConcurrentHandling(t *testing.T) { client := NewClient(testImpl, nil) var wg sync.WaitGroup for range 100 { - wg.Add(1) - go func() { - defer wg.Done() + wg.Go(func() { clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil) if err != nil { t.Errorf("Connect failed: %v", err) @@ -283,11 +288,17 @@ func TestStreamableConcurrentHandling(t *testing.T) { t.Errorf("CallTool failed: %v", err) return } - if got := int(res.StructuredContent.(map[string]any)["Count"].(float64)); got != i { - t.Errorf("got count %d, want %d", got, i) + number, ok := res.StructuredContent.(map[string]any)["Count"].(json.Number) + if !ok { + t.Errorf("count type = %T, want json.Number", res.StructuredContent.(map[string]any)["Count"]) + continue + } + got, err := number.Int64() + if err != nil || got != int64(i) { + t.Errorf("got count %d, %v; want %d", got, err, i) } } - }() + }) } wg.Wait() } @@ -334,10 +345,15 @@ func TestStreamableServerShutdown(t *testing.T) { opts.KeepAlive = 50 * time.Millisecond } client := NewClient(testImpl, &opts) + // Pin to 2025-11-25: this test relies on the legacy hanging GET + // (stateful case) or Ping-based keepalive (stateless case) for + // the client to notice the server shutdown. Both mechanisms are + // removed in 2026-07-28 (SEP-2575), so the new protocol would + // leave the client waiting forever. clientSession, err := client.Connect(ctx, &StreamableClientTransport{ Endpoint: httpServer.URL, MaxRetries: -1, // avoid slow tests during exponential retries - }, nil) + }, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatal(err) } @@ -367,6 +383,34 @@ func TestStreamableServerShutdown(t *testing.T) { } } +// TestStreamableStatelessKeepaliveRace verifies that there is no data race between +// ServerSession.startKeepalive and ServerSession.Close in stateless servers. +func TestStreamableStatelessKeepaliveRace(t *testing.T) { + ctx := context.Background() + server := NewServer(testImpl, &ServerOptions{KeepAlive: time.Hour}) + AddTool(server, &Tool{Name: "greet"}, sayHi) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true}, + ) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + for range 50 { + cs, err := NewClient(testImpl, nil).Connect(ctx, &StreamableClientTransport{ + Endpoint: httpServer.URL, + }, nil) + if err != nil { + t.Fatalf("NewClient() failed: %v", err) + } + _, _ = cs.CallTool(ctx, &CallToolParams{ + Name: "greet", + Arguments: map[string]any{"Name": "world"}, + }) + _ = cs.Close() + } +} + // TestClientReplay verifies that the client can recover from a mid-stream // network failure and receive replayed messages (if replay is configured). It // uses a proxy that is killed and restarted to simulate a recoverable network @@ -517,6 +561,92 @@ func testClientReplay(t *testing.T, test clientReplayTest) { } } +func TestStreamableServerDisconnect(t *testing.T) { + server := NewServer(testImpl, nil) + + // Test that client replayability allows the server to terminate incoming + // requests immediately, and have the client replay them. + + // testStream exercises stream resumption by interleaving stream termination + // with progress notifications. + testStream := func(ctx context.Context, session *ServerSession, extra *RequestExtra) { + // Close the stream before the first message. We should have sent an + // initial priming message already, so the client will be able to replay + extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: 10 * time.Millisecond}) + session.NotifyProgress(ctx, &ProgressNotificationParams{Message: "msg1"}) + time.Sleep(20 * time.Millisecond) + extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: 10 * time.Millisecond}) // Closing twice should still be supported. + session.NotifyProgress(ctx, &ProgressNotificationParams{Message: "msg2"}) + } + + AddTool(server, &Tool{Name: "disconnect"}, + func(ctx context.Context, req *CallToolRequest, args map[string]any) (*CallToolResult, map[string]any, error) { + testStream(ctx, req.Session, req.Extra) + return new(CallToolResult), nil, nil + }) + + server.AddPrompt(&Prompt{Name: "disconnect"}, func(ctx context.Context, req *GetPromptRequest) (*GetPromptResult, error) { + testStream(ctx, req.Session, req.Extra) + return nil, nil + }) + + tests := []struct { + name string + doCall func(context.Context, *ClientSession) error + }{ + { + "tool", + func(ctx context.Context, cs *ClientSession) error { + _, err := cs.CallTool(ctx, &CallToolParams{Name: "disconnect"}) + return err + }, + }, + { + "prompt", + func(ctx context.Context, cs *ClientSession) error { + _, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "disconnect"}) + return err + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + notifications := make(chan string, 2) + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{ + EventStore: NewMemoryEventStore(nil), // support replayability + }) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + client := NewClient(testImpl, &ClientOptions{ + ProgressNotificationHandler: func(ctx context.Context, req *ProgressNotificationClientRequest) { + notifications <- req.Params.Message + }, + }) + clientSession, err := client.Connect(ctx, &StreamableClientTransport{ + Endpoint: httpServer.URL, + }, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("client.Connect() failed: %v", err) + } + defer clientSession.Close() + + if err = test.doCall(ctx, clientSession); err != nil { + t.Fatalf("CallTool failed: %v", err) + } + + got := readNotifications(t, ctx, notifications, 2) + want := []string{"msg1", "msg2"} + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("got unexpected notifications (-want +got):\n%s", diff) + } + }) + } +} + func TestServerTransportCleanup(t *testing.T) { nClient := 3 @@ -540,7 +670,10 @@ func TestServerTransportCleanup(t *testing.T) { handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil) handler.onTransportDeletion = func(sessionID string) { - chans[sessionID] <- struct{}{} + mu.Lock() + ch := chans[sessionID] + mu.Unlock() + ch <- struct{}{} } httpServer := httptest.NewServer(mustNotPanic(t, handler)) @@ -565,14 +698,23 @@ func TestServerTransportCleanup(t *testing.T) { } } client.AddReceivingMiddleware(pingMiddleware) - clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil) + // Pin to 2025-11-25: this test counts server-side sessions, but + // the 2026-07-28 client probes server/discover first against a + // stateful HTTP server (which rejects the new protocol), opening + // an extra HTTP connection and therefore an extra session. Pinning + // to 2025-11-25 skips the discover probe. + clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatalf("client.Connect() failed: %v", err) } t.Cleanup(func() { _ = clientSession.Close() }) } - for _, ch := range chans { + mu.Lock() + channels := slices.Collect(maps.Values(chans)) + mu.Unlock() + + for _, ch := range channels { select { case <-ctx.Done(): t.Errorf("did not capture transport deletion event from all session in 10 seconds") @@ -594,14 +736,7 @@ func TestServerInitiatedSSE(t *testing.T) { notifications := make(chan string) server := NewServer(testImpl, nil) - opts := &StreamableHTTPOptions{ - // TODO(#583): for now, this is required for guaranteed message delivery. - // However, it shouldn't be necessary to use replay here, as we should be - // guaranteed that the standalone SSE stream is started by the time the - // client is connected. - EventStore: NewMemoryEventStore(nil), - } - httpServer := httptest.NewServer(mustNotPanic(t, NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, opts))) + httpServer := httptest.NewServer(mustNotPanic(t, NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil))) defer httpServer.Close() ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) @@ -659,6 +794,21 @@ func req(id int64, method string, params any) *jsonrpc.Request { return r } +func completeCallToolResult() *CallToolResult { + r := &CallToolResult{Content: []Content{}} + r.resultType = resultTypeComplete + return r +} + +// completeCallToolResultWithServerInfo returns a completeCallToolResult +// annotated with the auto-populated [MetaKeyServerInfo] entry that servers +// add to every new-protocol response. +func completeCallToolResultWithServerInfo(impl *Implementation) *CallToolResult { + r := completeCallToolResult() + r.Meta = Meta{MetaKeyServerInfo: impl} + return r +} + func resp(id int64, result any, err error) *jsonrpc.Response { return &jsonrpc.Response{ ID: jsonrpc2.Int64ID(id), @@ -673,13 +823,13 @@ func TestStreamableServerTransport(t *testing.T) { // requests. // Predefined steps, to avoid repetition below. - initReq := req(1, methodInitialize, &InitializeParams{}) + initReq := req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20250618}) initResp := resp(1, &InitializeResult{ Capabilities: &ServerCapabilities{ Logging: &LoggingCapabilities{}, Tools: &ToolCapabilities{ListChanged: true}, }, - ProtocolVersion: latestProtocolVersion, + ProtocolVersion: protocolVersion20250618, ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"}, }, nil) initializedMsg := req(0, notificationInitialized, &InitializedParams{}) @@ -696,11 +846,36 @@ func TestStreamableServerTransport(t *testing.T) { wantStatusCode: http.StatusAccepted, } + // Protocol version 2025-11-25 variants, for testing prime/close events (SEP-1699). + initReq20251125 := req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20251125}) + initResp20251125 := resp(1, &InitializeResult{ + Capabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Tools: &ToolCapabilities{ListChanged: true}, + }, + ProtocolVersion: protocolVersion20251125, + ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"}, + }, nil) + initialize20251125 := streamableRequest{ + method: "POST", + messages: []jsonrpc.Message{initReq20251125}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{initResp20251125}, + wantSessionID: true, + } + initialized20251125 := streamableRequest{ + method: "POST", + headers: http.Header{protocolVersionHeader: {protocolVersion20251125}}, + messages: []jsonrpc.Message{initializedMsg}, + wantStatusCode: http.StatusAccepted, + } + tests := []struct { - name string - replay bool // if set, use a MemoryEventStore to enable stream replay - tool func(*testing.T, context.Context, *ServerSession) - requests []streamableRequest // http requests + name string + replay bool // if set, use a MemoryEventStore to enable replay + tool func(*testing.T, context.Context, *CallToolRequest) // if set, called during execution + requests []streamableRequest + wantSessions int // number of sessions expected after the test }{ { name: "basic", @@ -714,6 +889,57 @@ func TestStreamableServerTransport(t *testing.T) { wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)}, }, }, + wantSessions: 1, + }, + { + name: "uninitialized", + requests: []streamableRequest{ + { + method: "POST", + messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusOK, + wantBodyContaining: "invalid during session initialization", + }, + }, + wantSessions: 0, + }, + { + name: "content type headers", + requests: []streamableRequest{ + initialize, + initialized, + { + // Request with incorrect Content-Type should be rejected. + method: "POST", + headers: http.Header{"Content-Type": {"text/plain"}}, + messages: []jsonrpc.Message{req(3, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusUnsupportedMediaType, + }, + { + // Request with empty Content-Type should be rejected. + method: "POST", + headers: http.Header{"Content-Type": {""}}, + messages: []jsonrpc.Message{req(4, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusUnsupportedMediaType, + }, + { + // Correct Content-Type should pass. + method: "POST", + headers: http.Header{"Content-Type": {"application/json"}}, + messages: []jsonrpc.Message{req(5, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(5, &CallToolResult{Content: []Content{}}, nil)}, + }, + { + // Correct Content-Type with parameters should pass. + method: "POST", + headers: http.Header{"Content-Type": {"application/json; charset=utf-8"}}, + messages: []jsonrpc.Message{req(5, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(5, &CallToolResult{Content: []Content{}}, nil)}, + }, + }, + wantSessions: 1, }, { name: "accept headers", @@ -742,12 +968,27 @@ func TestStreamableServerTransport(t *testing.T) { }, { method: "POST", - headers: http.Header{"Accept": {"text/*, application/*"}}, - messages: []jsonrpc.Message{req(4, "tools/call", &CallToolParams{Name: "tool"})}, + headers: http.Header{"Accept": {"application/json;charset=utf-8, text/event-stream"}}, + messages: []jsonrpc.Message{req(5, "tools/call", &CallToolParams{Name: "tool"})}, wantStatusCode: http.StatusOK, - wantMessages: []jsonrpc.Message{resp(4, &CallToolResult{Content: []Content{}}, nil)}, + wantMessages: []jsonrpc.Message{resp(5, &CallToolResult{Content: []Content{}}, nil)}, + }, + { + method: "POST", + headers: http.Header{"Accept": {"application/json;charset=utf-8", "text/event-stream"}}, + messages: []jsonrpc.Message{req(6, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(6, &CallToolResult{Content: []Content{}}, nil)}, + }, + { + method: "POST", + headers: http.Header{"Accept": {"text/*, application/*"}}, + messages: []jsonrpc.Message{req(7, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(7, &CallToolResult{Content: []Content{}}, nil)}, }, }, + wantSessions: 1, }, { name: "protocol version headers", @@ -763,6 +1004,7 @@ func TestStreamableServerTransport(t *testing.T) { wantSessionID: false, // could be true, but shouldn't matter }, }, + wantSessions: 1, }, { name: "batch rejected on 2025-06-18", @@ -782,6 +1024,7 @@ func TestStreamableServerTransport(t *testing.T) { wantBodyContaining: "batch", }, }, + wantSessions: 1, }, { name: "batch accepted on 2025-03-26", @@ -804,12 +1047,13 @@ func TestStreamableServerTransport(t *testing.T) { }, }, }, + wantSessions: 1, }, { name: "tool notification", - tool: func(t *testing.T, ctx context.Context, ss *ServerSession) { + tool: func(t *testing.T, ctx context.Context, req *CallToolRequest) { // Send an arbitrary notification. - if err := ss.NotifyProgress(ctx, &ProgressNotificationParams{}); err != nil { + if err := req.Session.NotifyProgress(ctx, &ProgressNotificationParams{}); err != nil { t.Errorf("Notify failed: %v", err) } }, @@ -828,12 +1072,13 @@ func TestStreamableServerTransport(t *testing.T) { }, }, }, + wantSessions: 1, }, { name: "tool upcall", - tool: func(t *testing.T, ctx context.Context, ss *ServerSession) { + tool: func(t *testing.T, ctx context.Context, req *CallToolRequest) { // Make an arbitrary call. - if _, err := ss.ListRoots(ctx, &ListRootsParams{}); err != nil { + if _, err := req.Session.ListRoots(ctx, &ListRootsParams{}); err != nil { t.Errorf("Call failed: %v", err) } }, @@ -860,26 +1105,29 @@ func TestStreamableServerTransport(t *testing.T) { }, }, }, + wantSessions: 1, }, { name: "background", // Enabling replay is necessary here because the standalone "GET" request - // is fully asynronous. Replay is needed to guarantee message delivery. + // is fully asynchronous. Replay is needed to guarantee message delivery. + // + // TODO(rfindley): this should no longer be necessary. replay: true, - tool: func(t *testing.T, _ context.Context, ss *ServerSession) { + tool: func(t *testing.T, _ context.Context, req *CallToolRequest) { // Perform operations on a background context, and ensure the client // receives it. ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) defer cancel() - if err := ss.NotifyProgress(ctx, &ProgressNotificationParams{}); err != nil { + if err := req.Session.NotifyProgress(ctx, &ProgressNotificationParams{}); err != nil { t.Errorf("Notify failed: %v", err) } // TODO(rfindley): finish implementing logging. // if err := ss.LoggingMessage(ctx, &LoggingMessageParams{}); err != nil { // t.Errorf("Logging failed: %v", err) // } - if _, err := ss.ListRoots(ctx, &ListRootsParams{}); err != nil { + if _, err := req.Session.ListRoots(ctx, &ListRootsParams{}); err != nil { t.Errorf("ListRoots failed: %v", err) } }, @@ -922,6 +1170,99 @@ func TestStreamableServerTransport(t *testing.T) { headers: map[string][]string{"Accept": nil}, }, }, + wantSessions: 0, // session deleted + }, + { + name: "no priming message on old protocol", + replay: true, + requests: []streamableRequest{ + initialize, + initialized, + { + method: "POST", + messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)}, + wantBodyNotContaining: "prime", + }, + }, + wantSessions: 1, + }, + { + name: "no close message on old protocol", + replay: true, + tool: func(t *testing.T, _ context.Context, req *CallToolRequest) { + req.Extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: time.Millisecond}) + }, + requests: []streamableRequest{ + initialize, + initialized, + { + method: "POST", + messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)}, + wantBodyNotContaining: "close", + }, + }, + wantSessions: 1, + }, + { + name: "priming message on 2025-11-25", + replay: true, + requests: []streamableRequest{ + initialize20251125, + initialized20251125, + { + method: "POST", + headers: http.Header{protocolVersionHeader: {protocolVersion20251125}}, + messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)}, + wantBodyContaining: "prime", + }, + }, + wantSessions: 1, + }, + { + name: "close message on 2025-11-25", + replay: true, + tool: func(t *testing.T, _ context.Context, req *CallToolRequest) { + req.Extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: time.Millisecond}) + }, + requests: []streamableRequest{ + initialize20251125, + initialized20251125, + { + method: "POST", + headers: http.Header{protocolVersionHeader: {protocolVersion20251125}}, + messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)}, + wantBodyContaining: "close", + wantBodyNotContaining: "result", + }, + }, + wantSessions: 1, + }, + { + name: "no close message", + replay: true, + tool: func(t *testing.T, _ context.Context, req *CallToolRequest) { + req.Extra.CloseSSEStream(CloseSSEStreamArgs{}) + }, + requests: []streamableRequest{ + initialize, + initialized, + { + method: "POST", + messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)}, + wantBodyNotContaining: "close", + }, + }, + wantSessions: 1, }, { name: "errors", @@ -948,11 +1289,12 @@ func TestStreamableServerTransport(t *testing.T) { method: "POST", messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})}, wantStatusCode: http.StatusOK, - wantMessages: []jsonrpc.Message{resp(2, nil, &jsonrpc2.WireError{ + wantMessages: []jsonrpc.Message{resp(2, nil, &jsonrpc.Error{ Message: `method "tools/call" is invalid during session initialization`, })}, }, }, + wantSessions: 0, }, } @@ -965,7 +1307,7 @@ func TestStreamableServerTransport(t *testing.T) { &Tool{Name: "tool", InputSchema: &jsonschema.Schema{Type: "object"}}, func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { if test.tool != nil { - test.tool(t, ctx, req.Session) + test.tool(t, ctx, req) } return &CallToolResult{}, nil }) @@ -979,6 +1321,9 @@ func TestStreamableServerTransport(t *testing.T) { defer handler.closeAll() testStreamableHandler(t, handler, test.requests) + if got := len(slices.Collect(server.Sessions())); got != test.wantSessions { + t.Errorf("after test, got %d sessions, want %d", got, test.wantSessions) + } }) } } @@ -1034,9 +1379,7 @@ func testStreamableHandler(t *testing.T, handler http.Handler, requests []stream ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() + wg.Go(func() { for m := range out { if req, ok := m.(*jsonrpc.Request); ok && req.IsCall() { @@ -1056,7 +1399,7 @@ func testStreamableHandler(t *testing.T, handler http.Handler, requests []stream cancel() } } - }() + }) gotSessionID, gotStatusCode, gotBody, err := request.do(ctx, httpServer.URL, sessionID.Load().(string), out) @@ -1074,11 +1417,14 @@ func testStreamableHandler(t *testing.T, handler http.Handler, requests []stream } wg.Wait() - if request.wantBodyContaining != "" { + if request.wantBodyContaining != "" || request.wantBodyNotContaining != "" { body := string(gotBody) - if !strings.Contains(body, request.wantBodyContaining) { + if request.wantBodyContaining != "" && !strings.Contains(body, request.wantBodyContaining) { t.Errorf("body does not contain %q:\n%s", request.wantBodyContaining, body) } + if request.wantBodyNotContaining != "" && strings.Contains(body, request.wantBodyNotContaining) { + t.Errorf("body contains %q:\n%s", request.wantBodyNotContaining, body) + } } else { transform := cmpopts.AcyclicTransformer("jsonrpcid", func(id jsonrpc.ID) any { return id.Raw() }) if diff := cmp.Diff(request.wantMessages, got, transform); diff != "" { @@ -1091,11 +1437,9 @@ func testStreamableHandler(t *testing.T, handler http.Handler, requests []stream var wg sync.WaitGroup for i, request := range requests { if request.async || request.onRequest > 0 { - wg.Add(1) - go func() { - defer wg.Done() + wg.Go(func() { doStep(t, i, request) - }() + }) } else { doStep(t, i, request) } @@ -1130,11 +1474,12 @@ type streamableRequest struct { headers http.Header // additional headers to set, overlaid on top of the default headers messages []jsonrpc.Message // messages to send - closeAfter int // if nonzero, close after receiving this many messages - wantStatusCode int // expected status code - wantBodyContaining string // if set, expect the response body to contain this text; overrides wantMessages - wantMessages []jsonrpc.Message // expected messages to receive; ignored if wantBodyContaining is set - wantSessionID bool // whether or not a session ID is expected in the response + closeAfter int // if nonzero, close after receiving this many messages + wantStatusCode int // expected status code + wantBodyContaining string // if set, expect the response body to contain this text; overrides wantMessages + wantBodyNotContaining string // if set, a negative assertion on the body; overrides wantMessages + wantMessages []jsonrpc.Message // expected messages to receive; ignored if wantBodyContaining is set + wantSessionID bool // whether or not a session ID is expected in the response } // streamingRequest makes a request to the given streamable server with the @@ -1183,10 +1528,16 @@ func (s streamableRequest) do(ctx context.Context, serverURL, sessionID string, if sessionID != "" { req.Header.Set(sessionIDHeader, sessionID) } - req.Header.Set("Content-Type", "application/json") + if s.method == http.MethodPost { + req.Header.Set("Content-Type", "application/json") + } req.Header.Set("Accept", "application/json, text/event-stream") maps.Copy(req.Header, s.headers) + if req.Header.Get("Content-Type") == "" { + req.Header.Del("Content-Type") + } + resp, err := http.DefaultClient.Do(req) if err != nil { return "", 0, nil, fmt.Errorf("request failed: %v", err) @@ -1195,24 +1546,26 @@ func (s streamableRequest) do(ctx context.Context, serverURL, sessionID string, newSessionID := resp.Header.Get(sessionIDHeader) - contentType := resp.Header.Get("Content-Type") + contentType := baseMediaType(resp.Header.Get("Content-Type")) var respBody []byte - if strings.HasPrefix(contentType, "text/event-stream") { + if contentType == "text/event-stream" { r := readerInto{resp.Body, new(bytes.Buffer)} - for evt, err := range scanEvents(r) { + for evt, err := range scanEventsLimited(r, DefaultMaxEventSize) { if err != nil { return newSessionID, resp.StatusCode, nil, fmt.Errorf("reading events: %v", err) } - // TODO(rfindley): do we need to check evt.name? - // Does the MCP spec say anything about this? - msg, err := jsonrpc2.DecodeMessage(evt.Data) - if err != nil { - return newSessionID, resp.StatusCode, nil, fmt.Errorf("decoding message: %w", err) + if evt.Name == "" || evt.Name == "message" { // ordinary message + // TODO(rfindley): do we need to check evt.name? + // Does the MCP spec say anything about this? + msg, err := jsonrpc2.DecodeMessage(evt.Data) + if err != nil { + return newSessionID, resp.StatusCode, nil, fmt.Errorf("decoding message: %w", err) + } + out <- msg } - out <- msg } respBody = r.w.Bytes() - } else if strings.HasPrefix(contentType, "application/json") { + } else if contentType == "application/json" { data, err := io.ReadAll(resp.Body) if err != nil { return newSessionID, resp.StatusCode, nil, fmt.Errorf("reading json body: %w", err) @@ -1315,7 +1668,7 @@ func TestStreamableStateless(t *testing.T) { Logging: &LoggingCapabilities{}, Tools: &ToolCapabilities{ListChanged: true}, }, - ProtocolVersion: latestProtocolVersion, + ProtocolVersion: protocolVersion20251125, ServerInfo: &Implementation{Name: "test", Version: "v1.0.0"}, }, nil) // This version of sayHi expects @@ -1390,40 +1743,16 @@ func TestStreamableStateless(t *testing.T) { } } - sessionlessHandler := NewStreamableHTTPHandler(func(*http.Request) *Server { - // Return a stateless server which never assigns a session ID. - server := NewServer(testImpl, &ServerOptions{ - GetSessionID: func() string { return "" }, - }) - AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi) - return server - }, &StreamableHTTPOptions{ - Stateless: true, - }) - - // First, test the "sessionless" stateless mode, where there is no session ID. - t.Run("sessionless", func(t *testing.T) { - testStreamableHandler(t, sessionlessHandler, requests) - testClientCompatibility(t, sessionlessHandler) - }) - - // Next, test the default stateless mode, where session IDs are permitted. - // - // This can be used by tools to look up application state preserved across - // subsequent requests. - requests[0].wantSessionID = true // now expect a session ID for initialize statelessHandler := NewStreamableHTTPHandler(func(*http.Request) *Server { - // Return a server with default options which should assign a random session ID. server := NewServer(testImpl, nil) AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi) return server }, &StreamableHTTPOptions{ Stateless: true, }) - t.Run("stateless", func(t *testing.T) { - testStreamableHandler(t, statelessHandler, requests) - testClientCompatibility(t, sessionlessHandler) - }) + + testStreamableHandler(t, statelessHandler, requests) + testClientCompatibility(t, statelessHandler) } func textContent(t *testing.T, res *CallToolResult) string { @@ -1438,24 +1767,20 @@ func textContent(t *testing.T, res *CallToolResult) string { return text.Text } -func TestTokenInfo(t *testing.T) { - oldAuth := testAuth.Load() - defer testAuth.Store(oldAuth) - testAuth.Store(true) +func TestSessionHijackingPrevention(t *testing.T) { + // This test verifies that sessions bound to a user ID cannot be accessed + // by a different user (session hijacking prevention). ctx := context.Background() - // Create a server with a tool that returns TokenInfo. - tokenInfo := func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, any, error) { - return &CallToolResult{Content: []Content{&TextContent{Text: fmt.Sprintf("%v", req.Extra.TokenInfo)}}}, nil, nil - } server := NewServer(testImpl, nil) - AddTool(server, &Tool{Name: "tokenInfo", Description: "return token info"}, tokenInfo) - streamHandler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil) - verifier := func(context.Context, string, *http.Request) (*auth.TokenInfo, error) { + + // Use the bearer token directly as the user ID. This simulates how a real + // verifier might extract a user ID from a JWT "sub" claim or introspection. + verifier := func(_ context.Context, token string, _ *http.Request) (*auth.TokenInfo, error) { return &auth.TokenInfo{ - Scopes: []string{"scope"}, - // Expiration is far, far in the future. + Scopes: []string{"scope"}, + UserID: token, Expiration: time.Date(5000, 1, 2, 3, 4, 5, 0, time.UTC), }, nil } @@ -1463,27 +1788,58 @@ func TestTokenInfo(t *testing.T) { httpServer := httptest.NewServer(mustNotPanic(t, handler)) defer httpServer.Close() - transport := &StreamableClientTransport{Endpoint: httpServer.URL} - client := NewClient(testImpl, nil) - session, err := client.Connect(ctx, transport, nil) - if err != nil { - t.Fatalf("client.Connect() failed: %v", err) + // Helper to send a JSON-RPC request as a given user. + doRequest := func(msg jsonrpc.Message, sessionID, userID string) *http.Response { + t.Helper() + data, _ := jsonrpc2.EncodeMessage(msg) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, httpServer.URL, bytes.NewReader(data)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set("Authorization", "Bearer "+userID) + if sessionID != "" { + req.Header.Set(sessionIDHeader, sessionID) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + return resp } - defer session.Close() - res, err := session.CallTool(ctx, &CallToolParams{Name: "tokenInfo"}) - if err != nil { - t.Fatal(err) + // Create a session as user1. + initReq := &jsonrpc.Request{Method: "initialize", ID: jsonrpc2.Int64ID(1)} + initReq.Params, _ = json.Marshal(&InitializeParams{ + ProtocolVersion: protocolVersion20250618, + ClientInfo: &Implementation{Name: "test", Version: "1.0"}, + }) + resp := doRequest(initReq, "", "user1") + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + t.Fatalf("initialize failed with status %d: %s", resp.StatusCode, body) } - if len(res.Content) == 0 { - t.Fatal("missing content") + sessionID := resp.Header.Get(sessionIDHeader) + if sessionID == "" { + t.Fatal("no session ID in response") } - tc, ok := res.Content[0].(*TextContent) - if !ok { - t.Fatal("not TextContent") + + pingReq := &jsonrpc.Request{Method: "ping", ID: jsonrpc2.Int64ID(2)} + pingReq.Params, _ = json.Marshal(&PingParams{}) + + // Try to access the session as user2 - should fail. + resp2 := doRequest(pingReq, sessionID, "user2") + defer resp2.Body.Close() + if resp2.StatusCode != http.StatusForbidden { + body, _ := io.ReadAll(resp2.Body) + t.Errorf("expected status %d for user mismatch, got %d: %s", http.StatusForbidden, resp2.StatusCode, body) } - if g, w := tc.Text, "&{[scope] 5000-01-02 03:04:05 +0000 UTC map[]}"; g != w { - t.Errorf("got %q, want %q", g, w) + + // Access as original user1 should succeed. + resp3 := doRequest(pingReq, sessionID, "user1") + defer resp3.Body.Close() + if resp3.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp3.Body) + t.Errorf("expected status %d for matching user, got %d: %s", http.StatusOK, resp3.StatusCode, body) } } @@ -1525,7 +1881,9 @@ func TestStreamableGET(t *testing.T) { if err != nil { t.Fatal(err) } - if got, want := resp.StatusCode, http.StatusMethodNotAllowed; got != want { + // GET without session should return 400 Bad Request (not 405) because + // GET is a valid method - it just requires a session ID. + if got, want := resp.StatusCode, http.StatusBadRequest; got != want { t.Errorf("initial GET: got status %d, want %d", got, want) } defer resp.Body.Close() @@ -1573,6 +1931,643 @@ func TestStreamableGET(t *testing.T) { } } +// TestStreamableMcpHeaderValidation tests the server-side Mcp-Method, +// Mcp-Name, and Mcp-Param header validation through the full HTTP handler, +// as specified in SEP-2243. +func TestStreamableMcpHeaderValidation(t *testing.T) { + // Temporarily register the future version so the handler accepts it. + + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + server.AddTool( + &Tool{Name: "my-tool", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{}, nil + }) + server.AddTool( + &Tool{ + Name: "execute_sql", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "Region", + }, + "query": map[string]any{ + "type": "string", + }, + }, + }, + }, + func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{}, nil + }) + + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{ + Stateless: true, + }) + defer handler.closeAll() + + testMeta := Meta{ + MetaKeyProtocolVersion: minVersionForStandardHeaders, + MetaKeyClientInfo: map[string]any{"name": "testClient", "version": "v1.0.0"}, + MetaKeyClientCapabilities: map[string]any{}, + } + serverImpl := &Implementation{Name: "testServer", Version: "v1.0.0"} + + testStreamableHandler(t, handler, []streamableRequest{ + { + method: "POST", + headers: http.Header{ + protocolVersionHeader: {minVersionForStandardHeaders}, + methodHeader: {"tools/call"}, + nameHeader: {"my-tool"}, + }, + messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(2, completeCallToolResultWithServerInfo(serverImpl), nil)}, + }, + { + method: "POST", + headers: http.Header{ + protocolVersionHeader: {minVersionForStandardHeaders}, + methodHeader: {"prompts/get"}, + nameHeader: {"my-tool"}, + }, + messages: []jsonrpc.Message{req(3, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})}, + wantStatusCode: http.StatusBadRequest, + wantBodyContaining: "Mcp-Method header value", + }, + { + method: "POST", + headers: http.Header{ + protocolVersionHeader: {minVersionForStandardHeaders}, + methodHeader: {"tools/call"}, + nameHeader: {"wrong-tool"}, + }, + messages: []jsonrpc.Message{req(4, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})}, + wantStatusCode: http.StatusBadRequest, + wantBodyContaining: "Mcp-Name header value", + }, + { + method: "POST", + headers: http.Header{ + protocolVersionHeader: {minVersionForStandardHeaders}, + methodHeader: {"TOOLS/CALL"}, + nameHeader: {"my-tool"}, + }, + messages: []jsonrpc.Message{req(5, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})}, + wantStatusCode: http.StatusBadRequest, + wantBodyContaining: "Mcp-Method header value", + }, + { + method: "POST", + headers: http.Header{ + protocolVersionHeader: {minVersionForStandardHeaders}, + methodHeader: {"tools/call"}, + nameHeader: {"my-tool"}, + }, + messages: []jsonrpc.Message{req(6, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(6, completeCallToolResultWithServerInfo(serverImpl), nil)}, + }, + { + method: "POST", + headers: http.Header{ + protocolVersionHeader: {minVersionForStandardHeaders}, + methodHeader: {"tools/call"}, + nameHeader: {"execute_sql"}, + paramHeaderPrefix + "Region": {"us-west1"}, + }, + messages: []jsonrpc.Message{req(7, "tools/call", &CallToolParams{ + Meta: testMeta, + Name: "execute_sql", + Arguments: map[string]any{"region": "us-west1", "query": "SELECT 1"}, + })}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(7, completeCallToolResultWithServerInfo(serverImpl), nil)}, + }, + { + method: "POST", + headers: http.Header{ + protocolVersionHeader: {minVersionForStandardHeaders}, + methodHeader: {"tools/call"}, + nameHeader: {"execute_sql"}, + paramHeaderPrefix + "Region": {"eu-central1"}, + }, + messages: []jsonrpc.Message{req(8, "tools/call", &CallToolParams{ + Meta: testMeta, + Name: "execute_sql", + Arguments: map[string]any{"region": "us-west1"}, + })}, + wantStatusCode: http.StatusBadRequest, + wantBodyContaining: "header mismatch", + }, + { + method: "POST", + headers: http.Header{ + protocolVersionHeader: {minVersionForStandardHeaders}, + methodHeader: {"tools/call"}, + nameHeader: {"execute_sql"}, + }, + messages: []jsonrpc.Message{req(9, "tools/call", &CallToolParams{ + Meta: testMeta, + Name: "execute_sql", + Arguments: map[string]any{"region": "us-west1"}, + })}, + wantStatusCode: http.StatusBadRequest, + wantBodyContaining: "missing", + }, + }) +} + +// TestStreamableMcpHeaderValidationErrorFormat verifies that header +// validation errors return a JSON-RPC error with code -32020 and +// Content-Type application/json, per SEP-2243. +func TestStreamableMcpHeaderValidationErrorFormat(t *testing.T) { + + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + server.AddTool( + &Tool{Name: "my-tool", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{}, nil + }) + + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{ + Stateless: true, + }) + defer handler.closeAll() + + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + // Use the MCP client with a custom RoundTripper to inject a bad header. + var toolCallResp *http.Response + var toolCallRespBody []byte + + customClient := &http.Client{ + Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) { + var originalMethodHeader string + if req.Header.Get(methodHeader) == "tools/call" { + originalMethodHeader = req.Header.Get(methodHeader) + req.Header.Set(methodHeader, "wrong-method") + } + resp, err := http.DefaultTransport.RoundTrip(req) + if err == nil && originalMethodHeader == "tools/call" { + toolCallResp = resp + toolCallRespBody, _ = io.ReadAll(resp.Body) + } + return resp, err + }), + } + + clientTransport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + HTTPClient: customClient, + } + + client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, nil) + ctx := context.Background() + session, err := client.Connect(ctx, clientTransport, &ClientSessionOptions{ProtocolVersion: minVersionForStandardHeaders}) + if err != nil { + t.Fatal(err) + } + defer session.Close() + + _, err = session.CallTool(ctx, &CallToolParams{Name: "my-tool"}) + // We expect an error because the server should reject it. + if err == nil { + t.Error("CallTool succeeded unexpectedly") + } + + if toolCallResp == nil { + t.Fatal("no response captured") + } + + // Verify HTTP status code. + if toolCallResp.StatusCode != http.StatusBadRequest { + t.Errorf("status code = %d, want %d", toolCallResp.StatusCode, http.StatusBadRequest) + } + + // Verify Content-Type. + if baseMediaType(toolCallResp.Header.Get("Content-Type")) != "application/json" { + t.Errorf("Content-Type = %q, want %q", baseMediaType(toolCallResp.Header.Get("Content-Type")), "application/json") + } + + // Verify JSON-RPC error body contains error code -32020. + msg, err := jsonrpc2.DecodeMessage(toolCallRespBody) + if err != nil { + t.Fatalf("failed to decode message: %v", err) + } + resp, ok := msg.(*jsonrpc2.Response) + if !ok { + t.Fatalf("expected *jsonrpc2.Response, got %T", msg) + } + var wireErr *jsonrpc2.WireError + if !errors.As(resp.Error, &wireErr) { + t.Fatalf("expected *jsonrpc2.WireError, got %T", resp.Error) + } + if wireErr.Code != CodeHeaderMismatch { + t.Errorf("wireErr.Code = %d, want %d", wireErr.Code, CodeHeaderMismatch) + } + if !strings.Contains(wireErr.Message, "Mcp-Method header value") { + t.Errorf("wireErr.Message = %q, want it to contain %q", wireErr.Message, "Mcp-Method header value") + } +} + +// TestStreamableMcpHeaderVersionGating verifies that header validation +// is skipped for protocol versions older than minVersionForStandardHeaders. +func TestStreamableMcpHeaderVersionGating(t *testing.T) { + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + server.AddTool( + &Tool{Name: "my-tool", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{}, nil + }) + + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil) + defer handler.closeAll() + + initReq := req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20251125}) + initResp := resp(1, &InitializeResult{ + Capabilities: &ServerCapabilities{ + Logging: &LoggingCapabilities{}, + Tools: &ToolCapabilities{ListChanged: true}, + }, + ProtocolVersion: protocolVersion20251125, + ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"}, + }, nil) + + testStreamableHandler(t, handler, []streamableRequest{ + { + method: "POST", + messages: []jsonrpc.Message{initReq}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{initResp}, + wantSessionID: true, + }, + { + method: "POST", + headers: http.Header{protocolVersionHeader: {protocolVersion20251125}}, + messages: []jsonrpc.Message{req(0, notificationInitialized, &InitializedParams{})}, + wantStatusCode: http.StatusAccepted, + }, + // Requests with deliberately wrong MCP headers should still succeed + // because the protocol version is too old for validation. + { + method: "POST", + headers: http.Header{ + protocolVersionHeader: {protocolVersion20251125}, + methodHeader: {"wrong-method"}, + nameHeader: {"wrong-name"}, + }, + messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "my-tool"})}, + wantStatusCode: http.StatusOK, + wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)}, + }, + }) +} + +// TestStreamableParamHeadersClientSetsHeaders verifies that the client sets +// Mcp-Param-* headers on tool calls when the tool has x-mcp-header annotations. +func TestStreamableParamHeadersClientSetsHeaders(t *testing.T) { + + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + server.AddTool( + &Tool{ + Name: "execute_sql", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "Region", + }, + "query": map[string]any{ + "type": "string", + }, + }, + }, + }, + func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil + }) + + // Install a receiving middleware on the server that opts the tools/list + // response caching by setting a positive ttlMs. + server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + res, err := next(ctx, method, req) + if err == nil && method == methodListTools { + if r, ok := res.(*ListToolsResult); ok { + r.TTLMs = 60_000 + } + } + return res, err + } + }) + + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{ + Stateless: true, + }) + defer handler.closeAll() + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + var capturedHeaders http.Header + customClient := &http.Client{ + Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) { + if req.Header.Get(methodHeader) == "tools/call" { + capturedHeaders = req.Header.Clone() + } + return http.DefaultTransport.RoundTrip(req) + }), + } + + clientTransport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + HTTPClient: customClient, + } + + client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, nil) + ctx := context.Background() + session, err := client.Connect(ctx, clientTransport, &ClientSessionOptions{ProtocolVersion: minVersionForStandardHeaders}) + if err != nil { + t.Fatal(err) + } + defer session.Close() + + // ListTools to populate the tool cache (needed for param headers). + // Pass a non-nil params so the SEP-2575 per-request _meta triple is + // injected; injectMeta is a no-op when params is nil. + if _, err := session.ListTools(ctx, &ListToolsParams{}); err != nil { + t.Fatal(err) + } + + _, err = session.CallTool(ctx, &CallToolParams{ + Name: "execute_sql", + Arguments: map[string]any{"region": "us-west1", "query": "SELECT 1"}, + }) + if err != nil { + t.Fatal(err) + } + + if capturedHeaders == nil { + t.Fatal("no tool call headers captured") + } + if got := capturedHeaders.Get(methodHeader); got != "tools/call" { + t.Errorf("Mcp-Method = %q, want %q", got, "tools/call") + } + if got := capturedHeaders.Get(nameHeader); got != "execute_sql" { + t.Errorf("Mcp-Name = %q, want %q", got, "execute_sql") + } + if got := capturedHeaders.Get(paramHeaderPrefix + "Region"); got != "us-west1" { + t.Errorf("Mcp-Param-Region = %q, want %q", got, "us-west1") + } +} + +// TestStreamableFilterValidToolsIntegration verifies that AddTool rejects +// tools with invalid x-mcp-header annotations at registration time, and +// that valid tools are returned by ListTools. +func TestStreamableFilterValidToolsIntegration(t *testing.T) { + + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + noop := func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{}, nil + } + + // Valid tool with correct x-mcp-header annotation. + server.AddTool(&Tool{ + Name: "valid-tool", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{ + "type": "string", + "x-mcp-header": "Region", + }, + }, + }, + }, noop) + + // Invalid tool: x-mcp-header on an array type should panic. + func() { + defer func() { + r := recover() + if r == nil { + t.Fatal("AddTool with invalid x-mcp-header annotation did not panic") + } + }() + server.AddTool(&Tool{ + Name: "invalid-tool", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "items": map[string]any{ + "type": "array", + "x-mcp-header": "Items", + }, + }, + }, + }, noop) + }() + + // Tool with no x-mcp-header annotations (always valid). + server.AddTool(&Tool{ + Name: "plain-tool", + InputSchema: &jsonschema.Schema{Type: "object"}, + }, noop) + + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{ + Stateless: true, + }) + defer handler.closeAll() + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, nil) + ctx := context.Background() + session, err := client.Connect(ctx, &StreamableClientTransport{ + Endpoint: httpServer.URL, + }, &ClientSessionOptions{ProtocolVersion: minVersionForStandardHeaders}) + if err != nil { + t.Fatal(err) + } + defer session.Close() + + // Pass non-nil params so the SEP-2575 per-request _meta triple is + // injected; injectMeta is a no-op when params is nil. + result, err := session.ListTools(ctx, &ListToolsParams{}) + if err != nil { + t.Fatal(err) + } + + toolNames := make([]string, len(result.Tools)) + for i, tool := range result.Tools { + toolNames[i] = tool.Name + } + sort.Strings(toolNames) + + wantNames := []string{"plain-tool", "valid-tool"} + if !slices.Equal(toolNames, wantNames) { + t.Errorf("ListTools returned %v, want %v", toolNames, wantNames) + } +} + +// TestStreamable405AllowHeader verifies RFC 9110 §15.5.6 compliance: +// 405 Method Not Allowed responses MUST include an Allow header. +func TestStreamable405AllowHeader(t *testing.T) { + server := NewServer(testImpl, nil) + + tests := []struct { + name string + stateless bool + method string + wantStatus int + wantAllow string + }{ + { + name: "unsupported method stateful", + stateless: false, + method: "PUT", + wantStatus: http.StatusMethodNotAllowed, + wantAllow: "GET, POST, DELETE", + }, + { + name: "GET in stateless mode", + stateless: true, + method: "GET", + wantStatus: http.StatusMethodNotAllowed, + wantAllow: "POST", + }, + { + // In stateless mode, only POST is supported. + // DELETE returns 405, consistent with Allow: POST. + name: "DELETE without session stateless", + stateless: true, + method: "DELETE", + wantStatus: http.StatusMethodNotAllowed, + wantAllow: "POST", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + opts := &StreamableHTTPOptions{Stateless: tt.stateless} + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + req, err := http.NewRequest(tt.method, httpServer.URL, nil) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Accept", "application/json, text/event-stream") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + if got := resp.StatusCode; got != tt.wantStatus { + t.Errorf("status code: got %d, want %d", got, tt.wantStatus) + } + + allow := resp.Header.Get("Allow") + if allow != tt.wantAllow { + t.Errorf("Allow header: got %q, want %q", allow, tt.wantAllow) + } + }) + } +} + +// TestStreamableGETWithoutSession verifies that GET without session ID in stateful mode +// returns 400 Bad Request (not 405), since GET is a supported method that requires a session. +func TestStreamableGETWithoutSession(t *testing.T) { + server := NewServer(testImpl, nil) + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + req, err := http.NewRequest("GET", httpServer.URL, nil) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Accept", "text/event-stream") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + + // GET without session should return 400 Bad Request, not 405 Method Not Allowed, + // because GET is a valid method - it just requires a session ID. + if got, want := resp.StatusCode, http.StatusBadRequest; got != want { + t.Errorf("status code: got %d, want %d", got, want) + } + if got, want := strings.TrimSpace(string(body)), "Bad Request: GET requires an Mcp-Session-Id header"; got != want { + t.Errorf("body: got %q, want %q", got, want) + } +} + +func TestStreamableGETWithoutEventStreamAccept(t *testing.T) { + server := NewServer(testImpl, nil) + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + req, err := http.NewRequest("GET", httpServer.URL, nil) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Accept", "application/json") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + + if got, want := resp.StatusCode, http.StatusBadRequest; got != want { + t.Errorf("status code: got %d, want %d", got, want) + } + if got, want := strings.TrimSpace(string(body)), "Accept must contain 'text/event-stream' for GET requests"; got != want { + t.Errorf("body: got %q, want %q", got, want) + } +} + +func TestBaseMediaType(t *testing.T) { + tests := []struct { + name string + value string + want string + }{ + {name: "empty", want: ""}, + {name: "json", value: "application/json", want: "application/json"}, + {name: "json with params", value: "Application/JSON; charset=utf-8", want: "application/json"}, + {name: "event stream with params", value: "Text/Event-Stream; charset=utf-8", want: "text/event-stream"}, + {name: "invalid", value: "application/json; charset", want: ""}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := baseMediaType(test.value); got != test.want { + t.Errorf("baseMediaType(%q) = %q, want %q", test.value, got, test.want) + } + }) + } +} + func TestStreamableClientContextPropagation(t *testing.T) { type contextKey string const testKey = contextKey("test-key") @@ -1586,7 +2581,7 @@ func TestStreamableClientContextPropagation(t *testing.T) { switch req.Method { case "POST": w.Header().Set("Content-Type", "application/json") - w.Header().Set("Mcp-Session-Id", "test-session") + w.Header().Set(sessionIDHeader, "test-session") w.WriteHeader(http.StatusOK) w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2025-03-26","capabilities":{},"serverInfo":{"name":"test","version":"1.0"}}}`)) case "GET": @@ -1621,8 +2616,8 @@ func TestStreamableClientContextPropagation(t *testing.T) { cancel() select { case <-streamableConn.ctx.Done(): - case <-time.After(100 * time.Millisecond): - t.Error("Connection context was not cancelled when parent was cancelled") + t.Errorf("cancelling the connection context after successful connection broke the connection") + default: } } @@ -1650,8 +2645,12 @@ func TestStreamableSessionTimeout(t *testing.T) { defer httpServer.Close() // Connect a client to create a session. + // Pin to 2025-11-25: under 2026-07-28 the client probes server/discover + // first against a stateful HTTP server (which rejects the new + // protocol), opening an extra HTTP connection and therefore an extra + // session. Pinning to 2025-11-25 skips the discover probe. client := NewClient(testImpl, nil) - session, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil) + session, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) if err != nil { t.Fatalf("client.Connect() failed: %v", err) } @@ -1676,18 +2675,15 @@ func TestStreamableSessionTimeout(t *testing.T) { // Spin up two goroutines, each making a request every 10ms. These requests // should keep the server from timing out. var wg sync.WaitGroup - wg.Add(2) for range 2 { - go func() { - defer wg.Done() - + wg.Go(func() { for range 20 { if _, err := session.ListTools(ctx, nil); err != nil { t.Errorf("ListTools failed: %v", err) } time.Sleep(10 * time.Millisecond) } - }() + }) } wg.Wait() @@ -1728,3 +2724,1732 @@ func mustNotPanic(t *testing.T, h http.Handler) http.Handler { h.ServeHTTP(w, req) }) } + +// TestPingEventFiltering verifies that the streamable client correctly filters +// out SSE "ping" events, which are used for keep-alive but should not be +// treated as JSON-RPC messages. +// +// This test addresses issue #636: the client should skip non-"message" events +// according to the SSE specification. It tests the fix in processStream where +// events with evt.Name != "" && evt.Name != "message" are skipped. +func TestPingEventFiltering(t *testing.T) { + // This test verifies the low-level processStream filtering. + // We create a mock response with ping and message events. + + sseData := `event: ping +data: ping + +event: message +id: 1 +data: {"jsonrpc":"2.0","id":1,"result":{}} + +event: ping +data: keepalive + +` + + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseData))), + } + + // Create a minimal streamableClientConn for testing + ctx := t.Context() + + incoming := make(chan jsonrpc.Message, 10) + done := make(chan struct{}) + + conn := &streamableClientConn{ + ctx: ctx, + done: done, + incoming: incoming, + } + + // Create a test request + testReq := &jsonrpc.Request{ + ID: jsonrpc2.Int64ID(1), + Method: "test", + } + + // Process the stream + go conn.processStream(ctx, "test", resp, testReq) + + // Collect messages with timeout + var messages []jsonrpc.Message + timeout := time.After(1 * time.Second) + +collectLoop: + for { + select { + case msg := <-incoming: + messages = append(messages, msg) + // We expect only 1 message (the response), not the ping events + if len(messages) >= 1 { + break collectLoop + } + case <-timeout: + break collectLoop + } + } + + // Verify we only received the actual message, not the ping events + if len(messages) != 1 { + t.Errorf("got %d messages, want 1 (ping events should be filtered)", len(messages)) + for i, msg := range messages { + t.Logf("message %d: %T", i, msg) + } + } + + // Verify the message is the response + if len(messages) > 0 { + resp, ok := messages[0].(*jsonrpc.Response) + if !ok { + t.Errorf("first message is %T, want *jsonrpc.Response", messages[0]) + } else if resp.ID.Raw() != int64(1) { + t.Errorf("response ID is %v, want 1", resp.ID.Raw()) + } + } +} + +// TestProcessStreamPrimingEvent verifies that the streamable client correctly ignores +// SSE events with empty data buffers, which are used as priming events (e.g. SEP-1699). +func TestProcessStreamPrimingEvent(t *testing.T) { + // We create a mock response with a priming event (empty data, with an ID), + // followed by a normal event. + sseData := `id: 123 + +id: 124 +data: {"jsonrpc":"2.0","id":1,"result":{}} + +` + + ctx := t.Context() + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(sseData)), + } + + incoming := make(chan jsonrpc.Message, 10) + done := make(chan struct{}) + + conn := &streamableClientConn{ + ctx: ctx, + done: done, + incoming: incoming, + failed: make(chan struct{}), + logger: ensureLogger(nil), + } + + lastID, _, clientClosed := conn.processStream(ctx, "test", resp, nil) + + if clientClosed { + t.Fatalf("processStream was unexpectedly closed by client") + } + + if lastID != "124" { + t.Errorf("lastEventID = %q, want %q", lastID, "124") + } + + select { + case msg := <-incoming: + if res, ok := msg.(*jsonrpc.Response); !(ok && res.ID == jsonrpc2.Int64ID(1)) { + t.Errorf("got unexpected message: %v", msg) + } + default: + t.Errorf("expected a JSON-RPC message to be produced") + } +} + +// TestProcessStreamMaxEventSize verifies that a streamableClientConn honors its +// configured maxEventSize. +func TestProcessStreamMaxEventSize(t *testing.T) { + jsonrpcEventOfSize := func(padSize int) string { + msg := fmt.Sprintf(`{"jsonrpc":"2.0","id":1,"result":{"pad":%q}}`, strings.Repeat("A", padSize)) + return "data: " + msg + "\n\n" + } + + tests := []struct { + name string + maxEventSize int + event string + wantFail bool + }{ + { + name: "event over cap is rejected", + event: jsonrpcEventOfSize(4096), + maxEventSize: 1024, + wantFail: true, + }, + { + name: "event under cap is accepted", + event: jsonrpcEventOfSize(1024), + maxEventSize: 4096, + wantFail: false, + }, + { + name: "negative cap disables the limit", + event: jsonrpcEventOfSize(4096), + maxEventSize: -1, + wantFail: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := t.Context() + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(tt.event)), + } + conn := &streamableClientConn{ + ctx: ctx, + done: make(chan struct{}), + incoming: make(chan jsonrpc.Message, 10), + failed: make(chan struct{}), + logger: ensureLogger(nil), + maxEventSize: tt.maxEventSize, + } + conn.processStream(ctx, "test", resp, nil) + + err := conn.failure() + if tt.wantFail && err == nil { + t.Fatal("failure() = nil, want a non-nil error for an oversized event") + } + if !tt.wantFail && err != nil { + t.Fatalf("failure() = %v, want nil", err) + } + }) + } +} + +// TestScanEventsPingFiltering is a unit test for the low-level event scanning +// with ping events to verify scanEvents properly parses all event types. +func TestScanEventsPingFiltering(t *testing.T) { + // Create SSE stream with mixed events + sseData := `event: ping +data: ping + +event: message +data: {"jsonrpc":"2.0","method":"test","params":{}} + +event: ping +data: keepalive + +event: message +data: {"jsonrpc":"2.0","method":"test2","params":{}} + +` + + reader := strings.NewReader(sseData) + var events []Event + + // Scan all events + for evt, err := range scanEventsLimited(reader, DefaultMaxEventSize) { + if err != nil { + if err != io.EOF { + t.Fatalf("scanEvents error: %v", err) + } + break + } + events = append(events, evt) + } + + // Verify we got all 4 events + if len(events) != 4 { + t.Fatalf("got %d events, want 4", len(events)) + } + + // Verify event types + expectedNames := []string{"ping", "message", "ping", "message"} + for i, evt := range events { + if evt.Name != expectedNames[i] { + t.Errorf("event %d: got name %q, want %q", i, evt.Name, expectedNames[i]) + } + } + + // Verify that we can decode the message events but would fail on ping events + for i, evt := range events { + switch evt.Name { + case "message": + _, err := jsonrpc.DecodeMessage(evt.Data) + if err != nil { + t.Errorf("event %d: failed to decode message event: %v", i, err) + } + case "ping": + // Ping events have non-JSON data and should fail decoding + _, err := jsonrpc.DecodeMessage(evt.Data) + if err == nil { + t.Errorf("event %d: ping event unexpectedly decoded as valid JSON-RPC", i) + } + } + } +} + +func Test_ExportErrSessionMissing(t *testing.T) { + ctx := context.Background() + + // 1. Setup server + impl := &Implementation{Name: "test", Version: "1.0.0"} + server := NewServer(impl, nil) + handler := NewStreamableHTTPHandler(func(r *http.Request) *Server { return server }, nil) + ts := httptest.NewServer(handler) + defer ts.Close() + + // 2. Setup client + clientTransport := &StreamableClientTransport{ + Endpoint: ts.URL, + } + client := NewClient(impl, nil) + // Pin to 2025-11-25: under 2026-07-28 the client probes server/discover + // first against a stateful HTTP server (which rejects the new + // protocol), opening an extra HTTP connection and therefore an extra + // session. Pinning to 2025-11-25 skips the discover probe. + session, err := client.Connect(ctx, clientTransport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("Connect failed: %v", err) + } + defer session.Close() + + // 3. Manually invalidate session on server + handler.mu.Lock() + if len(handler.sessions) != 1 { + handler.mu.Unlock() + t.Fatalf("expected 1 session, got %d", len(handler.sessions)) + } + for id := range handler.sessions { + delete(handler.sessions, id) + } + handler.mu.Unlock() + + // 4. Try to call a tool (or any request) + _, err = session.ListTools(ctx, nil) + if err == nil { + t.Fatal("expected error, got nil") + } + + // 5. Verify it's ErrSessionMissing + if !errors.Is(err, ErrSessionMissing) { + t.Errorf("expected error to wrap ErrSessionMissing, got: %v", err) + } +} + +// TestStreamableLocalhostProtection verifies that DNS rebinding protection +// is automatically enabled for localhost servers. +func TestStreamableLocalhostProtection(t *testing.T) { + server := NewServer(testImpl, nil) + + tests := []struct { + name string + listenAddr string // Address to listen on + hostHeader string // Host header in request + disableProtection bool // DisableLocalhostProtection setting + wantStatus int + }{ + // Auto-enabled for localhost listeners (127.0.0.1). + { + name: "127.0.0.1 accepts 127.0.0.1", + listenAddr: "127.0.0.1:0", + hostHeader: "127.0.0.1:1234", + disableProtection: false, + wantStatus: http.StatusOK, + }, + { + name: "127.0.0.1 accepts localhost", + listenAddr: "127.0.0.1:0", + hostHeader: "localhost:1234", + disableProtection: false, + wantStatus: http.StatusOK, + }, + { + name: "127.0.0.1 rejects evil.com", + listenAddr: "127.0.0.1:0", + hostHeader: "evil.com", + disableProtection: false, + wantStatus: http.StatusForbidden, + }, + { + name: "127.0.0.1 rejects evil.com:80", + listenAddr: "127.0.0.1:0", + hostHeader: "evil.com:80", + disableProtection: false, + wantStatus: http.StatusForbidden, + }, + { + name: "127.0.0.1 rejects localhost.evil.com", + listenAddr: "127.0.0.1:0", + hostHeader: "localhost.evil.com", + disableProtection: false, + wantStatus: http.StatusForbidden, + }, + + // When listening on 0.0.0.0, requests arriving via localhost are still protected + // because LocalAddrContextKey returns the actual connection's local address. + // This is actually more secure - DNS rebinding attacks target localhost regardless + // of the listener configuration. + { + name: "0.0.0.0 via localhost rejects evil.com", + listenAddr: "0.0.0.0:0", + hostHeader: "evil.com", + disableProtection: false, + wantStatus: http.StatusForbidden, + }, + + // Explicit disable + { + name: "disabled accepts evil.com", + listenAddr: "127.0.0.1:0", + hostHeader: "evil.com", + disableProtection: true, + wantStatus: http.StatusOK, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + opts := &StreamableHTTPOptions{ + Stateless: true, // Simpler for testing + DisableLocalhostProtection: tt.disableProtection, + } + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts) + + listener, err := net.Listen("tcp", tt.listenAddr) + if err != nil { + t.Fatalf("Failed to listen on %s: %v", tt.listenAddr, err) + } + defer listener.Close() + + srv := &http.Server{Handler: handler} + go srv.Serve(listener) + defer srv.Close() + + reqReader := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`) + req, err := http.NewRequest("POST", fmt.Sprintf("http://%s", listener.Addr().String()), reqReader) + if err != nil { + t.Fatal(err) + } + req.Host = tt.hostHeader + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + if got := resp.StatusCode; got != tt.wantStatus { + t.Errorf("Status code: got %d, want %d", got, tt.wantStatus) + } + }) + } +} + +func TestStreamableOriginProtection(t *testing.T) { + server := NewServer(testImpl, nil) + + tests := []struct { + name string + protection *http.CrossOriginProtection + requestOrigin string + wantStatusCode int + }{ + { + name: "no protection with Origin header", + protection: nil, + requestOrigin: "https://example.com", + wantStatusCode: http.StatusOK, + }, + { + name: "custom protection with trusted origin and same Origin", + protection: func() *http.CrossOriginProtection { + p := http.NewCrossOriginProtection() + if err := p.AddTrustedOrigin("https://example.com"); err != nil { + t.Fatal(err) + } + return p + }(), + requestOrigin: "https://example.com", + wantStatusCode: http.StatusOK, + }, + { + name: "custom protection with trusted origin and different Origin", + protection: func() *http.CrossOriginProtection { + p := http.NewCrossOriginProtection() + if err := p.AddTrustedOrigin("https://example.com"); err != nil { + t.Fatal(err) + } + return p + }(), + requestOrigin: "https://malicious.com", + wantStatusCode: http.StatusForbidden, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + opts := &StreamableHTTPOptions{ + Stateless: true, // avoid session ID requirement + CrossOriginProtection: tt.protection, + } + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + reqReader := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`) + req, err := http.NewRequest(http.MethodPost, httpServer.URL, reqReader) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Origin", tt.requestOrigin) + req.Header.Set("Accept", "application/json, text/event-stream") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + if got := resp.StatusCode; got != tt.wantStatusCode { + t.Errorf("Status code: got %d, want %d", got, tt.wantStatusCode) + } + }) + } +} + +// TestStandaloneSSEEmitsCommentForHTTP2Flush is a regression test for the +// HTTP/2 reverse-proxy header-buffering bug. The standalone SSE GET stream +// must emit at least one body byte (an SSE comment) immediately after the +// response headers. Without a DATA frame, HTTP/2 reverse proxies (Envoy, +// Caddy, net/http/httputil, etc.) buffer the HEADERS frame indefinitely +// because they have nothing to coalesce it with — calling Flush() alone is +// not sufficient. +// +// SSE explicitly defines lines starting with ":" as comments that clients +// ignore, so this is behavior-preserving for the SSE protocol while +// producing the DATA frame that HTTP/2 proxies need. +func TestStandaloneSSEEmitsCommentForHTTP2Flush(t *testing.T) { + server := NewServer(testImpl, nil) + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + // Initialize a session so we have a valid Mcp-Session-Id to use on the + // standalone SSE GET. + initBody := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`) + initReq, err := http.NewRequestWithContext(ctx, http.MethodPost, httpServer.URL, initBody) + if err != nil { + t.Fatal(err) + } + initReq.Header.Set("Content-Type", "application/json") + initReq.Header.Set("Accept", "application/json, text/event-stream") + initResp, err := http.DefaultClient.Do(initReq) + if err != nil { + t.Fatal(err) + } + io.Copy(io.Discard, initResp.Body) + initResp.Body.Close() + sessionID := initResp.Header.Get(sessionIDHeader) + if sessionID == "" { + t.Fatalf("initialize response missing %s header", sessionIDHeader) + } + + // Open the standalone SSE stream. + getReq, err := http.NewRequestWithContext(ctx, http.MethodGet, httpServer.URL, nil) + if err != nil { + t.Fatal(err) + } + getReq.Header.Set("Accept", "text/event-stream") + getReq.Header.Set(sessionIDHeader, sessionID) + + resp, err := http.DefaultClient.Do(getReq) + if err != nil { + t.Fatalf("GET standalone SSE: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d; want 200", resp.StatusCode) + } + if got := resp.Header.Get("Content-Type"); !strings.HasPrefix(got, "text/event-stream") { + t.Fatalf("Content-Type = %q; want text/event-stream", got) + } + + // Read the first chunk. With the fix, an SSE comment (": ok\n\n") is + // written immediately. Without the fix, the body has no bytes to read + // until the server pushes an event, which won't happen in this test. + type readResult struct { + n int + err error + buf []byte + } + ch := make(chan readResult, 1) + go func() { + buf := make([]byte, 64) + n, err := resp.Body.Read(buf) + ch <- readResult{n: n, err: err, buf: buf} + }() + + select { + case r := <-ch: + if r.err != nil && r.err != io.EOF { + t.Fatalf("reading first SSE chunk: %v", r.err) + } + if r.n == 0 { + t.Fatal("first SSE chunk was empty; expected an SSE comment to flush HTTP/2 proxy HEADERS frame") + } + got := string(r.buf[:r.n]) + // SSE spec: lines starting with ":" are comments, ignored by clients. + // We don't pin the exact comment text; just require the first byte to + // be the SSE comment marker so HTTP/2 proxies have a DATA frame. + if !strings.HasPrefix(got, ":") { + t.Errorf("first SSE chunk = %q; want it to start with %q (SSE comment marker, so HTTP/2 proxies receive a DATA frame and forward HEADERS)", got, ":") + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for first SSE bytes; the standalone SSE stream must emit a DATA frame immediately so HTTP/2 reverse proxies don't buffer the HEADERS frame") + } +} + +// newProtocolBody builds a raw JSON body for a tools/call request that +// carries the >= 2026-07-28 per-request _meta fields. +func newProtocolBody(t *testing.T, toolName string, args any) []byte { + t.Helper() + rawArgs, err := json.Marshal(args) + if err != nil { + t.Fatal(err) + } + body, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "new-proto-client", "version": "9.9"}, + MetaKeyClientCapabilities: map[string]any{"sampling": map[string]any{}}, + }, + "name": toolName, + "arguments": json.RawMessage(rawArgs), + }, + }) + if err != nil { + t.Fatal(err) + } + return body +} + +func TestEphemeralConnectOpts(t *testing.T) { + mkReq := func(body []byte) *http.Request { + r := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body)) + r.Header.Set("Content-Type", "application/json") + return r + } + + h := &StreamableHTTPHandler{opts: StreamableHTTPOptions{}} + + oldProtocolBody, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": map[string]any{"name": "x", "arguments": map[string]any{}}, + }) + if err != nil { + t.Fatal(err) + } + initializeBody, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": methodInitialize, + "params": map[string]any{"protocolVersion": protocolVersion20250618}, + }) + if err != nil { + t.Fatal(err) + } + + tests := []struct { + name string + body []byte + wantUsesNew bool + wantInitializeParams bool + wantInitializedParams bool + }{ + { + name: "new-protocol request: no synthetic state", + body: newProtocolBody(t, "x", struct{}{}), + wantUsesNew: true, + wantInitializeParams: false, + wantInitializedParams: false, + }, + { + name: "old-protocol request: synthetic state populated", + body: oldProtocolBody, + wantUsesNew: false, + wantInitializeParams: true, + wantInitializedParams: true, + }, + { + name: "initialize request: no synthetic InitializeParams", + body: initializeBody, + wantUsesNew: false, + wantInitializeParams: false, + wantInitializedParams: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := mkReq(tt.body) + var pver string + if tt.wantUsesNew { + pver = protocolVersion20260728 + } else { + pver = protocolVersion20250326 + } + req.Header.Set(protocolVersionHeader, pver) + req = req.WithContext(context.WithValue(req.Context(), protocolVersionContextKey{}, pver)) + info, err := h.ephemeralConnectOpts(req) + if err != nil { + t.Fatal(err) + } + if info.usesNewProtocol != tt.wantUsesNew { + t.Errorf("usesNewProtocol = %v, want %v", info.usesNewProtocol, tt.wantUsesNew) + } + if got := info.opts.State.InitializeParams != nil; got != tt.wantInitializeParams { + t.Errorf("InitializeParams non-nil = %v, want %v (value = %+v)", + got, tt.wantInitializeParams, info.opts.State.InitializeParams) + } + if got := info.opts.State.InitializedParams != nil; got != tt.wantInitializedParams { + t.Errorf("InitializedParams non-nil = %v, want %v (value = %+v)", + got, tt.wantInitializedParams, info.opts.State.InitializedParams) + } + }) + } +} + +// statelessHandlerCapture builds a stateless server with a single tool whose +// handler captures everything we want to assert about the per-request view of +// the session and the new-protocol accessors. +type statelessHandlerCapture struct { + mu sync.Mutex + sessionInitParams *InitializeParams + reqProtocolVersion string + reqClientInfo *Implementation + reqClientCapabilities *ClientCapabilities +} + +func TestStreamableStateless_NewProtocolSession_NoFakeInit(t *testing.T) { + // SEP-2575: the MCP-Protocol-Version header is mandatory for new-protocol + // requests and must be a supported version. The 2026-07-28 version is + // not yet in the global list, so register it for the duration of the test. + + capture := &statelessHandlerCapture{} + mcpServer := NewServer(testImpl, nil) + AddTool(mcpServer, &Tool{Name: "capture", Description: "captures request info"}, + func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) { + capture.mu.Lock() + defer capture.mu.Unlock() + capture.sessionInitParams = req.Session.InitializeParams() + capture.reqProtocolVersion = req.ProtocolVersion() + capture.reqClientInfo = req.ClientInfo() + capture.reqClientCapabilities = req.ClientCapabilities() + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return mcpServer }, + &StreamableHTTPOptions{Stateless: true}, + ) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + body := newProtocolBody(t, "capture", struct{}{}) + httpReq, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "application/json, text/event-stream") + httpReq.Header.Set(protocolVersionHeader, protocolVersion20260728) + // >= 2026-07-28 also requires the Mcp-Method and Mcp-Name standard + // headers (see streamable_headers.go). + httpReq.Header.Set(methodHeader, "tools/call") + httpReq.Header.Set(nameHeader, "capture") + + resp, err := http.DefaultClient.Do(httpReq) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + respBody, _ := io.ReadAll(resp.Body) + t.Fatalf("status = %d, want 200; body = %s", resp.StatusCode, respBody) + } + + capture.mu.Lock() + defer capture.mu.Unlock() + if capture.sessionInitParams == nil { + t.Errorf("Session.InitializeParams() is nil, want populated initializeParams for new-protocol session") + } else { + if got, want := capture.sessionInitParams.ProtocolVersion, protocolVersion20260728; got != want { + t.Errorf("Session.InitializeParams().ProtocolVersion = %q, want %q", got, want) + } + if got, want := capture.sessionInitParams.ClientInfo.Name, "new-proto-client"; got != want { + t.Errorf("Session.InitializeParams().ClientInfo.Name = %q, want %q", got, want) + } + } + if got, want := capture.reqProtocolVersion, protocolVersion20260728; got != want { + t.Errorf("req.ProtocolVersion() = %q, want %q", got, want) + } + if capture.reqClientInfo == nil || capture.reqClientInfo.Name != "new-proto-client" { + t.Errorf("req.ClientInfo() = %+v, want Name=new-proto-client", capture.reqClientInfo) + } + if capture.reqClientCapabilities == nil || capture.reqClientCapabilities.Sampling == nil { + t.Errorf("req.ClientCapabilities() = %+v, want non-nil Sampling", capture.reqClientCapabilities) + } +} + +// TestStreamableStateful_RejectsNewProtocol verifies that a stateful HTTP +// server rejects requests carrying _meta.protocolVersion (i.e. >= 2026-07-28 +// requests) with HTTP 400. The new protocol is +// supported on HTTP only when StreamableHTTPOptions.Stateless=true. +func TestStreamableStateful_RejectsNewProtocol(t *testing.T) { + // Make 2026-07-28 a "known" version so that the request reaches servePOST + // (otherwise the early header validation at ServeHTTP rejects it). + + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "noop"}, + func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + // Initialize a legacy session first. + initBody := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`) + initReq, err := http.NewRequest(http.MethodPost, httpServer.URL, initBody) + if err != nil { + t.Fatal(err) + } + initReq.Header.Set("Content-Type", "application/json") + initReq.Header.Set("Accept", "application/json, text/event-stream") + initResp, err := http.DefaultClient.Do(initReq) + if err != nil { + t.Fatal(err) + } + io.Copy(io.Discard, initResp.Body) + initResp.Body.Close() + sessionID := initResp.Header.Get(sessionIDHeader) + if sessionID == "" { + t.Fatalf("initialize response missing %s header", sessionIDHeader) + } + + // Drive the existing session with a new-protocol request whose header and + // body agree. The cross-check passes; the stateful-rejection check fires. + body := newProtocolBody(t, "noop", struct{}{}) + req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(sessionIDHeader, sessionID) + req.Header.Set(protocolVersionHeader, protocolVersion20260728) + req.Header.Set(methodHeader, "tools/call") + req.Header.Set(nameHeader, "noop") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("status = %d, want 400; body = %s", resp.StatusCode, respBody) + } + msg, err := jsonrpc.DecodeMessage(respBody) + if err != nil { + t.Fatalf("decoding response body as JSON-RPC: %v; body = %s", err, respBody) + } + jresp, ok := msg.(*jsonrpc.Response) + if !ok || jresp.Error == nil { + t.Fatalf("response is not a JSON-RPC error: %s", respBody) + } + var jerr *jsonrpc.Error + if !errors.As(jresp.Error, &jerr) { + t.Fatalf("response error is not *jsonrpc.Error: %v", jresp.Error) + } + if jerr.Code != CodeUnsupportedProtocolVersion { + t.Errorf("error code = %d, want %d (CodeUnsupportedProtocolVersion)", jerr.Code, CodeUnsupportedProtocolVersion) + } + var data UnsupportedProtocolVersionData + if err := json.Unmarshal(jerr.Data, &data); err != nil { + t.Fatalf("decoding error data: %v; data = %s", err, jerr.Data) + } + if data.Requested != protocolVersion20260728 { + t.Errorf("data.Requested = %q, want %q", data.Requested, protocolVersion20260728) + } + if slices.Contains(data.Supported, protocolVersion20260728) { + t.Errorf("data.Supported = %v, must not contain the rejected version %q", + data.Supported, protocolVersion20260728) + } + if len(data.Supported) == 0 { + t.Errorf("data.Supported is empty; expected at least one legacy version") + } +} + +// TestStreamableStateful_RejectsNewProtocol_LegacyPlainText verifies that +// MCPGODEBUG=plaintextstatefulrejection=1 restores the pre-v1.8.0 plain-text +// 400 body for the stateful/new-protocol rejection. +func TestStreamableStateful_RejectsNewProtocol_LegacyPlainText(t *testing.T) { + prev := plaintextstatefulrejection + plaintextstatefulrejection = "1" + t.Cleanup(func() { plaintextstatefulrejection = prev }) + + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "noop"}, + func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + initBody := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`) + initReq, err := http.NewRequest(http.MethodPost, httpServer.URL, initBody) + if err != nil { + t.Fatal(err) + } + initReq.Header.Set("Content-Type", "application/json") + initReq.Header.Set("Accept", "application/json, text/event-stream") + initResp, err := http.DefaultClient.Do(initReq) + if err != nil { + t.Fatal(err) + } + io.Copy(io.Discard, initResp.Body) + initResp.Body.Close() + sessionID := initResp.Header.Get(sessionIDHeader) + if sessionID == "" { + t.Fatalf("initialize response missing %s header", sessionIDHeader) + } + + body := newProtocolBody(t, "noop", struct{}{}) + req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(sessionIDHeader, sessionID) + req.Header.Set(protocolVersionHeader, protocolVersion20260728) + req.Header.Set(methodHeader, "tools/call") + req.Header.Set(nameHeader, "noop") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("status = %d, want 400; body = %s", resp.StatusCode, respBody) + } + if ct := resp.Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/plain") { + t.Errorf("Content-Type = %q, want text/plain; body = %s", ct, respBody) + } + if !strings.Contains(string(respBody), "stateless") { + t.Errorf("body = %q, want a message mentioning 'stateless'", respBody) + } +} + +// TestCallCancellation_FastReturn verifies that a cancelled tool call +// returns as soon as its context is cancelled, even when the peer is slow +// to accept the follow-up notifications/cancelled POST. +func TestCallCancellation_FastReturn(t *testing.T) { + // stallDuration is the time the middleware holds a + // notifications/cancelled POST. It must be larger than + // notifyCancellationTimeout so the blocking branch clearly overshoots + // even after cushioning for scheduler jitter. + const stallDuration = notifyCancellationTimeout + 2*time.Second + + // callerDeadline is the deadline the client passes into CallTool. It + // must be short enough that any waiting on the cancel notification is + // obvious in the measured elapsed time. + const callerDeadline = 100 * time.Millisecond + + // slack accounts for scheduling jitter around the deadline. + const slack = 500 * time.Millisecond + + stallCancelNotify := func(h http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost && r.Body != nil { + body, _ := io.ReadAll(r.Body) + r.Body = io.NopCloser(bytes.NewReader(body)) + if bytes.Contains(body, []byte(notificationCancelled)) { + time.Sleep(stallDuration) + } + } + h.ServeHTTP(w, r) + }) + } + + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "slow"}, + func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) { + // Bound the handler above the caller deadline so the test + // always exercises the cancellation path, but not so long + // that a stuck cancellation makes the test slow. + select { + case <-ctx.Done(): + case <-time.After(2 * stallDuration): + } + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil) + httpServer := httptest.NewServer(stallCancelNotify(handler)) + t.Cleanup(httpServer.Close) + + client := NewClient(testImpl, nil) + cs, err := client.Connect(context.Background(), + &StreamableClientTransport{Endpoint: httpServer.URL}, nil) + if err != nil { + t.Fatalf("client connect: %v", err) + } + t.Cleanup(func() { cs.Close() }) + + callCtx, cancel := context.WithTimeout(context.Background(), callerDeadline) + defer cancel() + + start := time.Now() + _, err = cs.CallTool(callCtx, &CallToolParams{Name: "slow"}) + elapsed := time.Since(start) + + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("CallTool error = %v, want context.DeadlineExceeded", err) + } + // The caller should return within slack of its own deadline. + if elapsed > callerDeadline+slack { + t.Errorf("CallTool returned after %v; want <= %v (caller deadline %v + slack %v)", + elapsed, callerDeadline+slack, callerDeadline, slack) + } +} + +// TestStreamableStateless_AcceptsNewProtocol is the positive control: +// confirms that a stateless server still accepts new-protocol requests +// (the rejection in TestStreamableStateful_RejectsNewProtocol must not +// fire on Stateless: true). +func TestStreamableStateless_AcceptsNewProtocol(t *testing.T) { + + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "noop"}, + func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true}, + ) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + body := newProtocolBody(t, "noop", struct{}{}) + req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(protocolVersionHeader, protocolVersion20260728) + req.Header.Set(methodHeader, "tools/call") + req.Header.Set(nameHeader, "noop") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + respBody, _ := io.ReadAll(resp.Body) + t.Fatalf("status = %d, want 200; body = %s", resp.StatusCode, respBody) + } +} + +// TestStreamableStateless_NotificationMetaValidation checks that the SEP-2575 +// per-request `_meta` triple is required of calls only. NotificationParams +// declares `_meta` optional with no protocolVersion, so a notification that +// omits it is well-formed and must be accepted with 202. +func TestStreamableStateless_NotificationMetaValidation(t *testing.T) { + newProtocolMeta := map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "new-proto-client", "version": "9.9"}, + MetaKeyClientCapabilities: map[string]any{}, + } + + tests := []struct { + name string + message map[string]any + wantStatus int + }{ + { + name: "cancelled without meta", + message: map[string]any{ + "jsonrpc": "2.0", + "method": notificationCancelled, + "params": map[string]any{"requestID": 1, "reason": "context canceled"}, + }, + wantStatus: http.StatusAccepted, + }, + { + name: "progress without meta", + message: map[string]any{ + "jsonrpc": "2.0", + "method": notificationProgress, + "params": map[string]any{"progressToken": "t", "progress": 1}, + }, + wantStatus: http.StatusAccepted, + }, + { + name: "notification with matching meta", + message: map[string]any{ + "jsonrpc": "2.0", + "method": notificationCancelled, + "params": map[string]any{ + "_meta": newProtocolMeta, + "requestID": 1, + }, + }, + wantStatus: http.StatusAccepted, + }, + { + // A notification that volunteers a version is still held to the + // header-match rule, so relaxing the requirement does not open a + // hole for inconsistent messages. + name: "notification with mismatched meta", + message: map[string]any{ + "jsonrpc": "2.0", + "method": notificationCancelled, + "params": map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: "2025-06-18", + MetaKeyClientCapabilities: map[string]any{}, + }, + "requestID": 1, + }, + }, + wantStatus: http.StatusBadRequest, + }, + { + // Calls remain subject to the requirement. + name: "call without meta", + message: map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/list", + "params": map[string]any{}, + }, + wantStatus: http.StatusBadRequest, + }, + } + + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "noop"}, + func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true}, + ) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + body, err := json.Marshal(test.message) + if err != nil { + t.Fatal(err) + } + req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(protocolVersionHeader, protocolVersion20260728) + req.Header.Set(methodHeader, test.message["method"].(string)) + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode != test.wantStatus { + t.Fatalf("status = %d, want %d; body = %s", resp.StatusCode, test.wantStatus, respBody) + } + if test.wantStatus == http.StatusAccepted && len(respBody) > 0 { + t.Errorf("accepted notification returned body %q, want empty", respBody) + } + }) + } +} + +// TestStreamableClientUnsupportedVersionFallback exercises the full +// SEP-2575 fallback. The client requests protocolVersion20260728, which the +// server is configured to NOT advertise in supportedProtocolVersions for the +// duration of this test. The server therefore rejects the server/discover +// POST at the transport-level header validation with a plain HTTP 400 ("Bad +// Request: Unsupported protocol version ..."). The streamable client must +// recognize this body, keep the connection alive, and successfully complete +// the legacy initialize handshake. +func TestStreamableClientUnsupportedVersionFallback(t *testing.T) { + // Remove protocolVersion20260728 from the server-side supported list so + // the fallback path is exercised. + orig := supportedProtocolVersions + supportedProtocolVersions = slices.DeleteFunc(slices.Clone(orig), func(v string) bool { + return v == protocolVersion20260728 + }) + t.Cleanup(func() { supportedProtocolVersions = orig }) + + ctx := context.Background() + + server := NewServer(testImpl, nil) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + nil, + ) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + client := NewClient(testImpl, nil) + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer session.Close() + + ir := session.InitializeResult() + if ir == nil { + t.Fatal("InitializeResult is nil; expected legacy initialize to populate it") + } + // The fallback initialize explicitly requests protocolVersion20251125 + // (see client.go), so the server negotiates that version. + if ir.ProtocolVersion != protocolVersion20251125 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (legacy fallback)", + ir.ProtocolVersion, protocolVersion20251125) + } + + // Verify the session is fully usable after the fallback by issuing a + // real call against the server. + if err := session.Ping(ctx, nil); err != nil { + t.Errorf("Ping after fallback initialize: %v", err) + } +} + +// TestStreamableStateful_AcceptsDiscover verifies that a stateful HTTP server +// accepts a server/discover probe carrying MCP-Protocol-Version: 2026-07-28 +// (and the matching _meta.protocolVersion), instead of rejecting it with the +// "stateless required" 400. The SEP-2575 client flow has the client probing +// the server with the new protocol version to learn which versions are +// supported. +func TestStreamableStateful_AcceptsDiscover(t *testing.T) { + + server := NewServer(testImpl, nil) + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + body, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": methodDiscover, + "params": map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "new-proto-client", "version": "9.9"}, + MetaKeyClientCapabilities: map[string]any{"sampling": map[string]any{}}, + }, + }, + }) + if err != nil { + t.Fatal(err) + } + req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(protocolVersionHeader, protocolVersion20260728) + req.Header.Set(methodHeader, methodDiscover) + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200; body = %s", resp.StatusCode, respBody) + } + + // Parse the JSON-RPC response. The body may arrive as a bare JSON object + // or as a single SSE event depending on the Accept negotiation; both + // shapes are valid here. + jsonPayload := respBody + if i := bytes.Index(respBody, []byte("data: ")); i >= 0 { + jsonPayload = respBody[i+len("data: "):] + if j := bytes.IndexByte(jsonPayload, '\n'); j >= 0 { + jsonPayload = jsonPayload[:j] + } + } + var rpcResp struct { + Result *DiscoverResult `json:"result"` + Error *jsonrpc.Error `json:"error"` + } + if err := json.Unmarshal(jsonPayload, &rpcResp); err != nil { + t.Fatalf("unmarshal response %q: %v", respBody, err) + } + if rpcResp.Error != nil { + t.Fatalf("discover returned error: %+v (body = %s)", rpcResp.Error, respBody) + } + if rpcResp.Result == nil { + t.Fatalf("discover returned no result; body = %s", respBody) + } + if slices.Contains(rpcResp.Result.SupportedVersions, protocolVersion20260728) { + t.Errorf("DiscoverResult.SupportedVersions = %v, must not include %q on a stateful transport", + rpcResp.Result.SupportedVersions, protocolVersion20260728) + } + if len(rpcResp.Result.SupportedVersions) == 0 { + t.Errorf("DiscoverResult.SupportedVersions is empty; want at least one legacy version") + } +} + +// On a stateful StreamableHTTPHandler, a v1.7.0-style client always +// probes server/discover before falling back to initialize. That probe used +// to create a ServerSession that was registered in the handler's session +// map but never surfaced via Mcp-Session-Id, so the client could not DELETE +// it and it was leaked for the lifetime of the process. +// +// The fix in Server.discover skips populating InitializeParams when the +// transport cannot serve the new protocol, which lets the safety-net +// cleanup in serveStatefulPOST close the otherwise-unaddressable session. +func TestStreamableStateful_DiscoverDoesNotLeakSession(t *testing.T) { + ctx := context.Background() + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "noop", Description: "noop"}, + func(context.Context, *CallToolRequest, struct{}) (*CallToolResult, any, error) { + return &CallToolResult{}, nil, nil + }) + + // Stateful (default). SessionTimeout unset so nothing masks the leak + // via idle eviction. + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + countSessions := func() int { + n := 0 + for range server.Sessions() { + n++ + } + return n + } + + const clients = 3 + for i := 1; i <= clients; i++ { + c := NewClient(&Implementation{Name: "c", Version: "1.0.0"}, nil) + sess, err := c.Connect(ctx, &StreamableClientTransport{ + Endpoint: httpServer.URL, + DisableStandaloneSSE: true, + MaxRetries: -1, + }, nil) + if err != nil { + t.Fatalf("client %d Connect: %v", i, err) + } + if _, err := sess.ListTools(ctx, nil); err != nil { + t.Fatalf("client %d ListTools: %v", i, err) + } + // Exactly one addressable session should exist while this client + // is connected. Anything more means a discover-only session + // slipped through. + if got := countSessions(); got != 1 { + t.Errorf("client %d connected: sessions = %d, want 1", i, got) + } + if err := sess.Close(); err != nil { + t.Fatalf("client %d Close: %v", i, err) + } + if got := countSessions(); got != 0 { + t.Errorf("client %d closed: sessions = %d, want 0", i, got) + } + } + + if got := countSessions(); got != 0 { + t.Errorf("final leaked sessions = %d, want 0", got) + } +} + +// TestStreamableHTTP_E2E_DiscoverSuccess is a full end-to-end smoke test for +// SEP-2575 over the streamable HTTP transport. +func TestStreamableHTTP_E2E_DiscoverSuccess(t *testing.T) { + ctx := context.Background() + server := NewServer(&Implementation{Name: "e2e-server", Version: "v1"}, nil) + // Register a simple tool so we can prove the session is usable end-to-end. + AddTool(server, &Tool{Name: "echo", Description: "echoes its input"}, + func(_ context.Context, _ *CallToolRequest, args struct { + Msg string `json:"msg"` + }) (*CallToolResult, struct{}, error) { + return &CallToolResult{ + Content: []Content{&TextContent{Text: args.Msg}}, + }, struct{}{}, nil + }, + ) + + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true}, + ) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + client := NewClient(&Implementation{Name: "e2e-client", Version: "v1"}, nil) + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + cs, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer cs.Close() + + ir := cs.InitializeResult() + if ir == nil { + t.Fatal("InitializeResult is nil after Connect; discover should have populated it") + } + if ir.ProtocolVersion != protocolVersion20260728 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (negotiated via discover)", + ir.ProtocolVersion, protocolVersion20260728) + } + if ir.ServerInfo == nil || ir.ServerInfo.Name != "e2e-server" { + t.Errorf("InitializeResult.ServerInfo = %+v, want name=e2e-server", ir.ServerInfo) + } + + // Prove the session is fully usable: list tools and call one. + tools, err := cs.ListTools(ctx, nil) + if err != nil { + t.Fatalf("ListTools: %v", err) + } + if len(tools.Tools) != 1 || tools.Tools[0].Name != "echo" { + t.Errorf("ListTools = %+v, want one tool named 'echo'", tools.Tools) + } + res, err := cs.CallTool(ctx, &CallToolParams{ + Name: "echo", + Arguments: map[string]any{"msg": "hello"}, + }) + if err != nil { + t.Fatalf("CallTool: %v", err) + } + if len(res.Content) != 1 { + t.Fatalf("CallTool result content = %+v, want 1 entry", res.Content) + } + tc, ok := res.Content[0].(*TextContent) + if !ok || tc.Text != "hello" { + t.Errorf("CallTool result[0] = %+v, want TextContent{Text:\"hello\"}", res.Content[0]) + } +} + +// A POST whose call ID is already in flight on the same session must be rejected +// with a JSON-RPC error, must not overwrite the existing request-to-stream +// mapping, and must not publish the duplicate message to the session's incoming +// channel. +func TestStreamableServerRejectsDuplicateInFlightRequestID(t *testing.T) { + id := jsonrpc2.Int64ID(1) + conn := &streamableServerConn{ + logger: ensureLogger(nil), + incoming: make(chan jsonrpc.Message, 1), + done: make(chan struct{}), + streams: map[string]*stream{ + "existing": { + id: "existing", + logger: ensureLogger(nil), + requests: map[jsonrpc.ID]struct{}{id: {}}, + }, + }, + requestStreams: map[jsonrpc.ID]string{id: "existing"}, + } + + data, err := jsonrpc2.EncodeMessage(req(1, methodPing, &PingParams{})) + if err != nil { + t.Fatalf("EncodeMessage() error = %v", err) + } + ctx, cancel := context.WithTimeout(t.Context(), 20*time.Millisecond) + defer cancel() + httpReq := httptest.NewRequestWithContext(ctx, http.MethodPost, "/", bytes.NewReader(data)) + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "application/json, text/event-stream") + + rec := httptest.NewRecorder() + conn.servePOST(rec, httpReq) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d; body = %s", rec.Code, http.StatusBadRequest, rec.Body.String()) + } + msg, err := jsonrpc2.DecodeMessage(rec.Body.Bytes()) + if err != nil { + t.Fatalf("DecodeMessage() error = %v; body = %s", err, rec.Body.String()) + } + resp, ok := msg.(*jsonrpc.Response) + if !ok { + t.Fatalf("response type = %T, want *jsonrpc.Response", msg) + } + if got := resp.ID.Raw(); got != int64(1) { + t.Fatalf("response ID = %v, want 1", got) + } + var jerr *jsonrpc.Error + if !errors.As(resp.Error, &jerr) { + t.Fatalf("response error = %v, want *jsonrpc.Error", resp.Error) + } + if jerr.Code != jsonrpc.CodeInvalidRequest { + t.Fatalf("error code = %d, want %d", jerr.Code, jsonrpc.CodeInvalidRequest) + } + if got := conn.requestStreams[id]; got != "existing" { + t.Fatalf("requestStreams[%v] = %q, want %q", id, got, "existing") + } + select { + case msg := <-conn.incoming: + t.Fatalf("duplicate request was published to incoming: %v", msg) + default: + } +} + +// TestStreamableServerPreservesSingleRequestBatchResponse verifies that a +// single-request JSON-RPC batch is returned as a one-element JSON array. +func TestStreamableServerPreservesSingleRequestBatchResponse(t *testing.T) { + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{JSONResponse: true}, + ) + defer handler.closeAll() + + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + body := []byte(`[{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test-client","version":"1.0"}}}]`) + req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + data, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want %d; body = %s", resp.StatusCode, http.StatusOK, data) + } + + var messages []json.RawMessage + if err := json.Unmarshal(data, &messages); err != nil { + t.Fatalf("response is not a JSON array: %v; body = %s", err, data) + } + if len(messages) != 1 { + t.Fatalf("response contains %d messages, want 1; body = %s", len(messages), data) + } + if _, err := jsonrpc2.DecodeMessage(messages[0]); err != nil { + t.Fatalf("response item is not a JSON-RPC message: %v", err) + } +} + +// TestStreamableMaxRequestBodyBytes verifies that the streamable HTTP handler +// enforces StreamableHTTPOptions.MaxRequestBodyBytes on incoming request +// bodies. The limit must apply uniformly regardless of transfer encoding: +// Content-Length, Transfer-Encoding: chunked, or HTTP/2 (which has no +// Content-Length). +func TestStreamableMaxRequestBodyBytes(t *testing.T) { + // A minimally-valid JSON-RPC initialize request. The body-size check runs + // before parsing, so exact contents don't matter for the reject cases, + // but we use a real one for the "under limit" success case. + validInit := []byte(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"t","version":"v1"}}}`) + + // Build a payload of the requested size by padding a JSON object. The + // server doesn't need to parse it; the MaxBytesReader trips first. + oversizedBody := func(n int) []byte { + b := make([]byte, n) + b[0] = '{' + for i := 1; i < n-1; i++ { + b[i] = ' ' + } + b[n-1] = '}' + return b + } + + tests := []struct { + name string + limit int64 // value assigned to MaxRequestBodyBytes + body []byte + wantStatus int + wantStatusDesc string + }{ + { + name: "default limit rejects oversized", + limit: 0, // triggers DefaultMaxRequestBodyBytes (4 MiB) + body: oversizedBody(DefaultMaxRequestBodyBytes + 1), + wantStatus: http.StatusRequestEntityTooLarge, + wantStatusDesc: "413 for body > 4MiB default", + }, + { + name: "custom limit rejects oversized", + limit: 1024, + body: oversizedBody(2048), + wantStatus: http.StatusRequestEntityTooLarge, + wantStatusDesc: "413 for body > custom 1KiB limit", + }, + { + name: "chunked encoding does not bypass limit", + limit: 1024, + body: oversizedBody(4096), + wantStatus: http.StatusRequestEntityTooLarge, + wantStatusDesc: "413 for chunked body > limit (regression: Content-Length-only checks are bypassable)", + }, + { + name: "negative limit disables the check", + limit: -1, + body: oversizedBody(2 * DefaultMaxRequestBodyBytes), + wantStatus: http.StatusBadRequest, // body is not valid JSON-RPC; but the size check must not fire + wantStatusDesc: "not 413 when limit is disabled", + }, + { + name: "under limit succeeds", + limit: 64 * 1024, + body: validInit, + wantStatus: http.StatusOK, + wantStatusDesc: "200 for a valid initialize under the limit", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + server := NewServer(testImpl, nil) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{ + Stateless: true, + MaxRequestBodyBytes: tc.limit, + }, + ) + ts := httptest.NewServer(handler) + defer ts.Close() + + req, err := http.NewRequest(http.MethodPost, ts.URL, bytes.NewReader(tc.body)) + if err != nil { + t.Fatalf("NewRequest: %v", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + + resp, err := ts.Client().Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + + if resp.StatusCode != tc.wantStatus { + t.Fatalf("%s: status = %d, want %d (body=%q)", + tc.wantStatusDesc, resp.StatusCode, tc.wantStatus, string(body)) + } + }) + } +} + +// TestStreamableSupportedProtocolVersions_Header verifies that the +// MCP-Protocol-Version header is validated against +// [ServerOptions.SupportedProtocolVersions], and not merely against the +// versions the SDK knows about. +func TestStreamableSupportedProtocolVersions_Header(t *testing.T) { + const initBody = `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}` + + newServer := func(t *testing.T, stateless bool) *httptest.Server { + server := NewServer(testImpl, &ServerOptions{ + SupportedProtocolVersions: []string{protocolVersion20251125}, + }) + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: stateless}) + httpServer := httptest.NewServer(handler) + t.Cleanup(httpServer.Close) + return httpServer + } + post := func(t *testing.T, url, version, sessionID string) *http.Response { + t.Helper() + req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(initBody)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(protocolVersionHeader, version) + if sessionID != "" { + req.Header.Set(sessionIDHeader, sessionID) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + return resp + } + + for _, stateless := range []bool{false, true} { + for _, test := range []struct { + version string + wantStatus int + }{ + {protocolVersion20251125, http.StatusOK}, + // Supported by the SDK, but excluded by the server. + {protocolVersion20250618, http.StatusBadRequest}, + // Unknown to the SDK: rejected before any server is consulted. + {"1999-01-01", http.StatusBadRequest}, + } { + t.Run(fmt.Sprintf("stateless=%v/%s", stateless, test.version), func(t *testing.T) { + resp := post(t, newServer(t, stateless).URL, test.version, "") + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode != test.wantStatus { + t.Errorf("status = %d, want %d; body = %s", resp.StatusCode, test.wantStatus, body) + } + }) + } + } + + t.Run("existing session", func(t *testing.T) { + httpServer := newServer(t, false) + resp := post(t, httpServer.URL, protocolVersion20251125, "") + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + sessionID := resp.Header.Get(sessionIDHeader) + if sessionID == "" { + t.Fatalf("initialize response missing %s header", sessionIDHeader) + } + + resp = post(t, httpServer.URL, protocolVersion20250618, sessionID) + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("status = %d, want %d; body = %s", resp.StatusCode, http.StatusBadRequest, body) + } + }) +} diff --git a/mcp/testdata/conformance/server/discover.txtar b/mcp/testdata/conformance/server/discover.txtar new file mode 100644 index 0000000..9159563 --- /dev/null +++ b/mcp/testdata/conformance/server/discover.txtar @@ -0,0 +1,44 @@ +The client sends a discover request carrying the per-request _meta +(protocolVersion, clientInfo, clientCapabilities), and the server +replies with the protocol versions it supports, its capabilities, and its +identity via `_meta.io.modelcontextprotocol/serverInfo`. + +-- client -- +{ + "jsonrpc": "2.0", + "id": 1, + "method": "server/discover", + "params": { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": { "name": "ExampleClient", "version": "1.0.0" }, + "io.modelcontextprotocol/clientCapabilities": {} + } + } +} +-- server -- +{ + "jsonrpc": "2.0", + "id": 1, + "result": { + "resultType": "complete", + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "testServer", + "version": "v1.0.0" + } + }, + "ttlMs": 0, + "cacheScope": "public", + "supportedVersions": [ + "2026-07-28", + "2025-11-25", + "2025-06-18", + "2025-03-26", + "2024-11-05" + ], + "capabilities": { + "logging": {} + } + } +} diff --git a/mcp/testdata/conformance/server/lifecycle.txtar b/mcp/testdata/conformance/server/lifecycle.txtar index 0a8cf34..652aa7f 100644 --- a/mcp/testdata/conformance/server/lifecycle.txtar +++ b/mcp/testdata/conformance/server/lifecycle.txtar @@ -53,6 +53,8 @@ See also modelcontextprotocol/go-sdk#225. "jsonrpc": "2.0", "id": 2, "result": { + "ttlMs": 0, + "cacheScope": "public", "tools": [] } } @@ -60,6 +62,8 @@ See also modelcontextprotocol/go-sdk#225. "jsonrpc": "2.0", "id": 3, "result": { + "ttlMs": 0, + "cacheScope": "public", "tools": [] } } diff --git a/mcp/testdata/conformance/server/mrtr.txtar b/mcp/testdata/conformance/server/mrtr.txtar new file mode 100644 index 0000000..046a79b --- /dev/null +++ b/mcp/testdata/conformance/server/mrtr.txtar @@ -0,0 +1,94 @@ +Check the multi-round-trip request pattern under protocol version +2026-07-28. The server's tool handler returns an +InputRequiredResult on the first call, requesting elicitation input from the +client and providing an opaque RequestState token. The client fulfills the +input request and retries the tool call with the responses and the echoed +RequestState; the handler then returns the final result. + +-- tools -- +confirmThenGreet +-- client -- +{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": { "name": "ExampleClient", "version": "1.0.0" }, + "io.modelcontextprotocol/clientCapabilities": { "elicitation": {} } + }, + "name": "confirmThenGreet", + "arguments": {} + } +} +{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": { "name": "ExampleClient", "version": "1.0.0" }, + "io.modelcontextprotocol/clientCapabilities": { "elicitation": {} } + }, + "name": "confirmThenGreet", + "arguments": {}, + "requestState": "step=1", + "inputResponses": { + "who": { "action": "accept", "content": { "name": "MCP Go" } } + } + } +} +-- server -- +{ + "jsonrpc": "2.0", + "id": 1, + "result": { + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "testServer", + "version": "v1.0.0" + } + }, + "content": null, + "requestState": "step=1", + "resultType": "input_required", + "inputRequests": { + "who": { + "method": "elicitation/create", + "params": { + "mode": "form", + "message": "What is your name?", + "requestedSchema": { + "type": "object", + "properties": { + "name": { + "type": "string" + } + } + } + } + } + } + } +} +{ + "jsonrpc": "2.0", + "id": 2, + "result": { + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "testServer", + "version": "v1.0.0" + } + }, + "content": [ + { + "type": "text", + "text": "hello MCP Go" + } + ], + "resultType": "complete" + } +} diff --git a/mcp/testdata/conformance/server/prompts.txtar b/mcp/testdata/conformance/server/prompts.txtar index fdaf793..2d8e168 100644 --- a/mcp/testdata/conformance/server/prompts.txtar +++ b/mcp/testdata/conformance/server/prompts.txtar @@ -45,6 +45,8 @@ code_review "jsonrpc": "2.0", "id": 2, "result": { + "ttlMs": 0, + "cacheScope": "public", "tools": [] } } @@ -52,6 +54,8 @@ code_review "jsonrpc": "2.0", "id": 4, "result": { + "ttlMs": 0, + "cacheScope": "public", "prompts": [ { "arguments": [ diff --git a/mcp/testdata/conformance/server/resources.txtar b/mcp/testdata/conformance/server/resources.txtar index 314817b..1c3ed80 100644 --- a/mcp/testdata/conformance/server/resources.txtar +++ b/mcp/testdata/conformance/server/resources.txtar @@ -67,6 +67,8 @@ info.txt "jsonrpc": "2.0", "id": 2, "result": { + "ttlMs": 0, + "cacheScope": "public", "resources": [ { "mimeType": "text/plain", @@ -85,6 +87,8 @@ info.txt "jsonrpc": "2.0", "id": 3, "result": { + "ttlMs": 0, + "cacheScope": "public", "contents": [ { "uri": "embedded:info", @@ -103,6 +107,8 @@ info.txt "jsonrpc": "2.0", "id": 3, "result": { + "ttlMs": 0, + "cacheScope": "public", "contents": [ { "uri": "file:///info.txt", diff --git a/mcp/testdata/conformance/server/spec-sep-973-additional-metadata.txtar b/mcp/testdata/conformance/server/spec-sep-973-additional-metadata.txtar new file mode 100644 index 0000000..356be18 --- /dev/null +++ b/mcp/testdata/conformance/server/spec-sep-973-additional-metadata.txtar @@ -0,0 +1,222 @@ +Check behavior of server with Icons and websiteUrl metadata added as part of SEP-973 specification. + +check modelcontextprotocol/go-sdk/issues/552 for more details. + +Checks following: +- Test setting websiteUrl, icons for mcp.Implementation +- Test setting icons for mcp.Prompt +- Test setting icons for mcp.Tool +- Test setting icons for mcp.Resource +- Test if tool call returning ResourceLink with icons works as expected + + +-- tools -- +greetWithIcon +contentTool + +-- prompts -- +code_reviewWithIcon + +-- resources -- +infoWithIcon + +-- client -- +{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": { "name": "ExampleClient", "version": "1.0.0" } + } +} +{ "jsonrpc": "2.0", "method": "notifications/initialized" } +{ "jsonrpc": "2.0", "id": 2, "method": "tools/list" } +{ "jsonrpc": "2.0", "id": 3, "method": "resources/list" } +{ "jsonrpc": "2.0", "id": 4, "method": "prompts/list" } +{ + "jsonrpc": "2.0", + "id": 5, + "method": "tools/call", + "params": { + "name": "contentTool", + "arguments": { + "In": "resource1" + } + } +} + +-- server -- +{ + "jsonrpc": "2.0", + "id": 1, + "result": { + "capabilities": { + "logging": {}, + "prompts": { + "listChanged": true + }, + "resources": { + "listChanged": true + }, + "tools": { + "listChanged": true + } + }, + "protocolVersion": "2025-06-18", + "serverInfo": { + "name": "testServer", + "version": "v1.0.0", + "websiteUrl": "https://github.com/modelcontextprotocol/go-sdk", + "icons": [ + { + "src": "foobar", + "mimeType": "image/png", + "sizes": [ + "48x48", + "96x96" + ], + "theme": "light" + } + ] + } + } +} +{ + "jsonrpc": "2.0", + "id": 2, + "result": { + "ttlMs": 0, + "cacheScope": "public", + "tools": [ + { + "description": "return resourceLink content with Icon", + "inputSchema": { + "type": "object", + "properties": { + "In": { + "type": "string", + "description": "the input" + } + }, + "required": [ + "In" + ], + "additionalProperties": false + }, + "name": "contentTool", + "title": "contentTool" + }, + { + "description": "say hi", + "inputSchema": { + "type": "object", + "properties": { + "Name": { + "type": "string" + } + }, + "required": [ + "Name" + ], + "additionalProperties": false + }, + "name": "greetWithIcon", + "icons": [ + { + "src": "foobar", + "mimeType": "image/png", + "sizes": [ + "48x48", + "96x96" + ], + "theme": "light" + } + ] + } + ] + } +} +{ + "jsonrpc": "2.0", + "id": 3, + "result": { + "ttlMs": 0, + "cacheScope": "public", + "resources": [ + { + "mimeType": "text/plain", + "name": "info", + "uri": "embedded:info", + "icons": [ + { + "src": "foobar", + "mimeType": "image/png", + "sizes": [ + "48x48", + "96x96" + ], + "theme": "light" + } + ] + } + ] + } +} +{ + "jsonrpc": "2.0", + "id": 4, + "result": { + "ttlMs": 0, + "cacheScope": "public", + "prompts": [ + { + "arguments": [ + { + "name": "Code", + "required": true + } + ], + "description": "do a code review", + "name": "code_review", + "icons": [ + { + "src": "foobar", + "mimeType": "image/png", + "sizes": [ + "48x48", + "96x96" + ], + "theme": "light" + } + ] + } + ] + } +} +{ + "jsonrpc": "2.0", + "id": 5, + "result": { + "content": [ + { + "type": "resource_link", + "mimeType": "text/plain", + "uri": "https://example.com/resource/resource1", + "name": "Example Resource Link with Icons", + "icons": [ + { + "src": "foobar", + "mimeType": "image/png", + "sizes": [ + "48x48", + "96x96" + ], + "theme": "light" + } + ] + } + ] + } +} diff --git a/mcp/testdata/conformance/server/tools.txtar b/mcp/testdata/conformance/server/tools.txtar index b582dda..ba6fd33 100644 --- a/mcp/testdata/conformance/server/tools.txtar +++ b/mcp/testdata/conformance/server/tools.txtar @@ -64,19 +64,21 @@ inc "jsonrpc": "2.0", "id": 2, "result": { + "ttlMs": 0, + "cacheScope": "public", "tools": [ { "description": "say hi", "inputSchema": { "type": "object", - "required": [ - "Name" - ], "properties": { "Name": { "type": "string" } }, + "required": [ + "Name" + ], "additionalProperties": false }, "name": "greet" @@ -95,70 +97,70 @@ inc "name": "inc", "outputSchema": { "type": "object", - "required": [ - "y" - ], "properties": { "y": { "type": "integer" } }, + "required": [ + "y" + ], "additionalProperties": false } }, { "inputSchema": { "type": "object", - "required": [ - "In" - ], "properties": { "In": { "type": "string", "description": "the input" } }, + "required": [ + "In" + ], "additionalProperties": false }, "name": "structured", "outputSchema": { "type": "object", - "required": [ - "Out" - ], "properties": { "Out": { "type": "string", "description": "the output" } }, + "required": [ + "Out" + ], "additionalProperties": false } }, { "inputSchema": { "type": "object", - "required": [ - "Now" - ], "properties": { "Now": { "type": "string" } }, + "required": [ + "Now" + ], "additionalProperties": false }, "name": "tomorrow", "outputSchema": { "type": "object", - "required": [ - "Tomorrow" - ], "properties": { "Tomorrow": { "type": "string" } }, + "required": [ + "Tomorrow" + ], "additionalProperties": false } } @@ -169,6 +171,8 @@ inc "jsonrpc": "2.0", "id": 3, "result": { + "ttlMs": 0, + "cacheScope": "public", "resources": [] } } @@ -176,6 +180,8 @@ inc "jsonrpc": "2.0", "id": 4, "result": { + "ttlMs": 0, + "cacheScope": "public", "prompts": [] } } @@ -222,9 +228,14 @@ inc { "jsonrpc": "2.0", "id": 8, - "error": { - "code": -32602, - "message": "invalid params: validating \"arguments\": validating root: required: missing properties: [\"In\"]" + "result": { + "content": [ + { + "type": "text", + "text": "validating \"arguments\": validating root: required: missing properties: [\"In\"]" + } + ], + "isError": true } } { @@ -245,9 +256,14 @@ inc { "jsonrpc": "2.0", "id": 10, - "error": { - "code": -32602, - "message": "invalid params: validating \"arguments\": validating root: required: missing properties: [\"Name\"]" + "result": { + "content": [ + { + "type": "text", + "text": "validating \"arguments\": validating root: required: missing properties: [\"Name\"]" + } + ], + "isError": true } } { diff --git a/mcp/testdata/conformance/server/version-latest.txtar b/mcp/testdata/conformance/server/version-latest.txtar index 7531767..77d1de2 100644 --- a/mcp/testdata/conformance/server/version-latest.txtar +++ b/mcp/testdata/conformance/server/version-latest.txtar @@ -20,7 +20,7 @@ response with its latest supported version. "capabilities": { "logging": {} }, - "protocolVersion": "2025-06-18", + "protocolVersion": "2025-11-25", "serverInfo": { "name": "testServer", "version": "v1.0.0" diff --git a/mcp/tool.go b/mcp/tool.go index 12b02b7..11ad100 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -8,8 +8,10 @@ import ( "context" "encoding/json" "fmt" + "strings" "github.com/google/jsonschema-go/jsonschema" + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" ) // A ToolHandler handles a call to tools/call. @@ -41,7 +43,7 @@ type ToolHandler func(context.Context, *CallToolRequest) (*CallToolResult, error // - If the Out type is not the empty interface [any], it provides the // default output schema for the tool (which again may be overridden in // [AddTool]). -// - The Out value is used to populate result.StructuredOutput. +// - The Out value is used to populate result.StructuredContent. // - If [CallToolResult.Content] is unset, it is populated with the JSON // content of the output. // - An error result is treated as a tool error, rather than a protocol @@ -63,8 +65,14 @@ type serverTool struct { // applySchema validates whether data is valid JSON according to the provided // schema, after applying schema defaults. // -// Returns the JSON value augmented with defaults. -func applySchema(data json.RawMessage, resolved *jsonschema.Resolved) (json.RawMessage, error) { +// If forOutput is false, the data is treated as tool input: the schema's root +// type must be "object" and the value is unmarshaled into a map. +// +// If forOutput is true, the data is treated as tool output: the schema's root +// may be of any type (object, array, primitive, composition). +// +// Returns the JSON value, augmented with defaults where applicable. +func applySchema(data json.RawMessage, resolved *jsonschema.Resolved, forOutput bool) (json.RawMessage, error) { // TODO: use reflection to create the struct type to unmarshal into. // Separate validation from assignment. @@ -73,31 +81,114 @@ func applySchema(data json.RawMessage, resolved *jsonschema.Resolved) (json.RawM // This avoids inconsistent representation due to custom marshallers, such as // time.Time (issue #449). // - // Additionally, unmarshalling into a map ensures that the resulting JSON is + // For input, unmarshalling into a map ensures that the resulting JSON is // at least {}, even if data is empty. For example, arguments is technically // an optional property of callToolParams, and we still want to apply the // defaults in this case. // // TODO(rfindley): in which cases can resolved be nil? - if resolved != nil { + if resolved == nil { + return data, nil + } + + var unmarshaled any + if !forOutput { v := make(map[string]any) if len(data) > 0 { - if err := json.Unmarshal(data, &v); err != nil { + if err := internaljson.Unmarshal(data, &v); err != nil { return nil, fmt.Errorf("unmarshaling arguments: %w", err) } } - if err := resolved.ApplyDefaults(&v); err != nil { - return nil, fmt.Errorf("applying schema defaults:\n%w", err) - } - if err := resolved.Validate(&v); err != nil { - return nil, err - } - // We must re-marshal with the default values applied. - var err error - data, err = json.Marshal(v) - if err != nil { - return nil, fmt.Errorf("marshalling with defaults: %v", err) + unmarshaled = v + } else { + if len(data) > 0 { + if err := internaljson.Unmarshal(data, &unmarshaled); err != nil { + return nil, fmt.Errorf("unmarshaling output: %w", err) + } } } - return data, nil + + // Apply defaults only when the value is a map: jsonschema.Resolved.ApplyDefaults + // only operates on object properties. For object-rooted output schemas, + // coerce a nil result (from "null" or empty data) into {} so handlers that + // return a typed-nil map still validate. + appliedDefaults := false + if _, ok := unmarshaled.(map[string]any); ok { + if err := resolved.ApplyDefaults(&unmarshaled); err != nil { + return nil, fmt.Errorf("applying schema defaults:\n%w", err) + } + appliedDefaults = true + } else if forOutput && unmarshaled == nil && resolved.Schema().Type == "object" { + unmarshaled = make(map[string]any) + if err := resolved.ApplyDefaults(&unmarshaled); err != nil { + return nil, fmt.Errorf("applying schema defaults:\n%w", err) + } + appliedDefaults = true + } + + if err := resolved.Validate(&unmarshaled); err != nil { + return nil, err + } + + // Re-marshal only when defaults may have changed the value. + if !appliedDefaults { + return data, nil + } + out, err := json.Marshal(unmarshaled) + if err != nil { + return nil, fmt.Errorf("marshalling with defaults: %v", err) + } + return out, nil +} + +// isObjectJSON reports whether data is a JSON object (i.e., starts with '{' +// after any leading whitespace). Returns false for arrays, primitives, null, +// or empty input. +func isObjectJSON(data json.RawMessage) bool { + for _, b := range data { + switch b { + case ' ', '\t', '\n', '\r': + continue + case '{': + return true + default: + return false + } + } + return false +} + +// validateToolName checks whether name is a valid tool name, reporting a +// non-nil error if not. +func validateToolName(name string) error { + if name == "" { + return fmt.Errorf("tool name cannot be empty") + } + if len(name) > 128 { + return fmt.Errorf("tool name exceeds maximum length of 128 characters (current: %d)", len(name)) + } + // For consistency with other SDKs, report characters in the order the appear + // in the name. + var invalidChars []string + seen := make(map[rune]bool) + for _, r := range name { + if !validToolNameRune(r) { + if !seen[r] { + invalidChars = append(invalidChars, fmt.Sprintf("%q", string(r))) + seen[r] = true + } + } + } + if len(invalidChars) > 0 { + return fmt.Errorf("tool name contains invalid characters: %s", strings.Join(invalidChars, ", ")) + } + return nil +} + +// validToolNameRune reports whether r is valid within tool names. +func validToolNameRune(r rune) bool { + return (r >= 'a' && r <= 'z') || + (r >= 'A' && r <= 'Z') || + (r >= '0' && r <= '9') || + r == '_' || r == '-' || r == '.' } diff --git a/mcp/tool_example_test.go b/mcp/tool_example_test.go index 8efc3ee..610755b 100644 --- a/mcp/tool_example_test.go +++ b/mcp/tool_example_test.go @@ -253,6 +253,74 @@ func ExampleAddTool_complexSchema() { // weather types: [sun partly_cloudy clouds rain snow] } +// !+contenttypes + +func ExampleAddTool_contentTypes() { + chartPNG := []byte("\x89PNG\r\n\x1a\n") + summaryWAV := []byte("RIFF....WAVE") + + server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) + mcp.AddTool(server, &mcp.Tool{Name: "report"}, func(context.Context, *mcp.CallToolRequest, struct{}) (*mcp.CallToolResult, any, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "Q3 revenue by region."}, + &mcp.ImageContent{Data: chartPNG, MIMEType: "image/png"}, + &mcp.AudioContent{Data: summaryWAV, MIMEType: "audio/wav"}, + // An embedded resource inlines the contents, so the client + // needs no follow-up resources/read. Binary contents go in + // ResourceContents.Blob instead of Text. + &mcp.EmbeddedResource{ + Resource: &mcp.ResourceContents{ + URI: "file:///reports/q3.csv", + MIMEType: "text/csv", + Text: "region,revenue\nEMEA,42\n", + }, + }, + // A resource link only points at a resource. Prefer it when the + // client may not need the contents, since an embedded resource + // is transferred either way. + &mcp.ResourceLink{URI: "file:///reports/q3.pdf", MIMEType: "application/pdf"}, + }, + }, nil, nil + }) + + ctx := context.Background() + session, err := connect(ctx, server) + if err != nil { + log.Fatal(err) + } + defer session.Close() + + res, err := session.CallTool(ctx, &mcp.CallToolParams{Name: "report"}) + if err != nil { + log.Fatal(err) + } + + // Clients type-switch over the content they receive. + for _, content := range res.Content { + switch c := content.(type) { + case *mcp.TextContent: + fmt.Println("text:", c.Text) + case *mcp.ImageContent: + fmt.Printf("image: %s, %d bytes\n", c.MIMEType, len(c.Data)) + case *mcp.AudioContent: + fmt.Printf("audio: %s, %d bytes\n", c.MIMEType, len(c.Data)) + case *mcp.EmbeddedResource: + fmt.Printf("embedded resource: %s (%s)\n", c.Resource.URI, c.Resource.MIMEType) + case *mcp.ResourceLink: + fmt.Printf("resource link: %s (%s)\n", c.URI, c.MIMEType) + } + } + // Output: + // text: Q3 revenue by region. + // image: image/png, 8 bytes + // audio: audio/wav, 12 bytes + // embedded resource: file:///reports/q3.csv (text/csv) + // resource link: file:///reports/q3.pdf (application/pdf) +} + +// !-contenttypes + func connect(ctx context.Context, server *mcp.Server) (*mcp.ClientSession, error) { t1, t2 := mcp.NewInMemoryTransports() if _, err := server.Connect(ctx, t1, nil); err != nil { diff --git a/mcp/tool_test.go b/mcp/tool_test.go index c440fcc..7a82c96 100644 --- a/mcp/tool_test.go +++ b/mcp/tool_test.go @@ -14,7 +14,7 @@ import ( "testing" "github.com/google/jsonschema-go/jsonschema" - "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) func TestApplySchema(t *testing.T) { @@ -46,7 +46,7 @@ func TestApplySchema(t *testing.T) { {`{"x": 0}`, new(map[string]any), &map[string]any{"x": 0.0}}, } { raw := json.RawMessage(tt.data) - raw, err = applySchema(raw, resolved) + raw, err = applySchema(raw, resolved, false) if err != nil { t.Fatal(err) } @@ -59,14 +59,91 @@ func TestApplySchema(t *testing.T) { } } +func TestApplySchemaOutput(t *testing.T) { + // SEP-2106: when forOutput is true, the schema may have a non-object root. + for _, tc := range []struct { + name string + schema *jsonschema.Schema + data string + want string + wantErr string + }{ + { + name: "array root accepts array", + schema: &jsonschema.Schema{Type: "array", Items: &jsonschema.Schema{Type: "integer"}}, + data: `[1,2,3]`, + want: `[1,2,3]`, + }, + { + name: "array root rejects wrong element type", + schema: &jsonschema.Schema{Type: "array", Items: &jsonschema.Schema{Type: "integer"}}, + data: `[1,"two",3]`, + wantErr: `"integer"`, + }, + { + name: "primitive number root", + schema: &jsonschema.Schema{Type: "number"}, + data: `42`, + want: `42`, + }, + { + name: "primitive string root", + schema: &jsonschema.Schema{Type: "string"}, + data: `"hello"`, + want: `"hello"`, + }, + { + name: "object root with defaults", + schema: &jsonschema.Schema{Type: "object", Properties: map[string]*jsonschema.Schema{"x": {Type: "integer", Default: json.RawMessage("3")}}}, + data: `{}`, + want: `{"x":3}`, + }, + { + name: "empty object root with defaults", + schema: &jsonschema.Schema{Type: "object", Properties: map[string]*jsonschema.Schema{"x": {Type: "integer", Default: json.RawMessage("3")}}}, + data: ``, + want: `{"x":3}`, + }, + { + name: "object root accepts null as empty object", + schema: &jsonschema.Schema{Type: "object"}, + data: `null`, + want: `{}`, + }, + } { + t.Run(tc.name, func(t *testing.T) { + resolved, err := tc.schema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true}) + if err != nil { + t.Fatal(err) + } + got, err := applySchema(json.RawMessage(tc.data), resolved, true) + if tc.wantErr != "" { + if err == nil { + t.Fatalf("got %s, want error containing %q", got, tc.wantErr) + } + if !strings.Contains(err.Error(), tc.wantErr) { + t.Errorf("got error %q, want containing %q", err, tc.wantErr) + } + return + } + if err != nil { + t.Fatal(err) + } + if string(got) != tc.want { + t.Errorf("got %s, want %s", got, tc.want) + } + }) + } +} + func TestToolErrorHandling(t *testing.T) { // Construct server and add both tools at the top level server := NewServer(testImpl, nil) // Create a tool that returns a structured error structuredErrorHandler := func(ctx context.Context, req *CallToolRequest, args map[string]any) (*CallToolResult, any, error) { - return nil, nil, &jsonrpc2.WireError{ - Code: codeInvalidParams, + return nil, nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, Message: "internal server error", } } @@ -106,13 +183,13 @@ func TestToolErrorHandling(t *testing.T) { t.Fatal("expected error, got nil") } - var wireErr *jsonrpc2.WireError + var wireErr *jsonrpc.Error if !errors.As(err, &wireErr) { - t.Fatalf("expected WireError, got %[1]T: %[1]v", err) + t.Fatalf("expected jsonrpc.Error, got %[1]T: %[1]v", err) } - if wireErr.Code != codeInvalidParams { - t.Errorf("expected error code %d, got %d", codeInvalidParams, wireErr.Code) + if wireErr.Code != jsonrpc.CodeInvalidParams { + t.Errorf("expected error code %d, got %d", jsonrpc.CodeInvalidParams, wireErr.Code) } }) @@ -145,3 +222,53 @@ func TestToolErrorHandling(t *testing.T) { } }) } + +func TestValidateToolName(t *testing.T) { + t.Run("valid", func(t *testing.T) { + validTests := []struct { + label string + toolName string + }{ + {"simple alphanumeric names", "getUser"}, + {"names with underscores", "get_user_profile"}, + {"names with dashes", "user-profile-update"}, + {"names with dots", "admin.tools.list"}, + {"mixed character names", "DATA_EXPORT_v2.1"}, + {"single character names", "a"}, + {"128 character names", strings.Repeat("a", 128)}, + } + for _, test := range validTests { + t.Run(test.label, func(t *testing.T) { + if err := validateToolName(test.toolName); err != nil { + t.Errorf("validateToolName(%q) = %v, want nil", test.toolName, err) + } + }) + } + }) + + t.Run("invalid", func(t *testing.T) { + invalidTests := []struct { + label string + toolName string + wantErrContaining string + }{ + {"empty names", "", "tool name cannot be empty"}, + {"names longer than 128 characters", strings.Repeat("a", 129), "tool name exceeds maximum length of 128 characters (current: 129)"}, + {"names with spaces", "get user profile", `tool name contains invalid characters: " "`}, + {"names with commas", "get,user,profile", `tool name contains invalid characters: ","`}, + {"names with forward slashes", "user/profile/update", `tool name contains invalid characters: "/"`}, + {"names with other special chars", "user@domain.com", `tool name contains invalid characters: "@"`}, + {"names with multiple invalid chars", "user name@domain,com", `tool name contains invalid characters: " ", "@", ","`}, + {"names with unicode characters", "user-ñame", `tool name contains invalid characters: "ñ"`}, + } + for _, test := range invalidTests { + t.Run(test.label, func(t *testing.T) { + if err := validateToolName(test.toolName); err == nil || !strings.Contains(err.Error(), test.wantErrContaining) { + t.Errorf("validateToolName(%q) = %v, want error containing %q", test.toolName, err, test.wantErrContaining) + } + }) + } + + }) + +} diff --git a/mcp/transport.go b/mcp/transport.go index cacd65f..08d3b15 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -10,20 +10,40 @@ import ( "errors" "fmt" "io" - "log" + "log/slog" "net" "os" "sync" + "time" + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" - "github.com/modelcontextprotocol/go-sdk/internal/xcontext" + "github.com/modelcontextprotocol/go-sdk/internal/mcpgodebug" "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) +// notifyCancellationTimeout bounds the cancellation notification we send to +// the peer when the caller's context is cancelled. The notification is +// best-effort: a degraded connection (e.g. an OAuth flow that has been +// abandoned) must not be able to block the caller's return path or +// re-trigger expensive recovery on its behalf. See issue #882. +const notifyCancellationTimeout = 5 * time.Second + +// blockingcancelnotify, when set to "1" via MCPGODEBUG, restores the previous +// behavior of blocking the caller's return on delivery of the best-effort +// notifications/cancelled message (up to notifyCancellationTimeout). By +// default, the call is retired immediately and the notification is sent +// asynchronously so the caller returns as soon as its context is cancelled. +var blockingcancelnotify = mcpgodebug.Value("blockingcancelnotify") + // ErrConnectionClosed is returned when sending a message to a connection that // is closed or in the process of closing. var ErrConnectionClosed = errors.New("connection closed") +// ErrSessionMissing is returned when the session is known to not be present on +// the server. +var ErrSessionMissing = errors.New("session not found") + // A Transport is used to create a bidirectional connection between MCP client // and server. // @@ -36,6 +56,19 @@ type Transport interface { Connect(ctx context.Context) (Connection, error) } +// ProtocolVersionSupporter is an optional capability that a [Transport] may +// implement to declare which MCP protocol versions it can serve. +// +// [Server.Connect] consults this interface to filter the +// list of versions advertised in server/discover responses. Transports that +// do not implement this interface are assumed to support every protocol +// version known to the SDK. +type ProtocolVersionSupporter interface { + // SupportsProtocolVersion reports whether the transport can serve + // requests using the given protocol version. + SupportsProtocolVersion(version string) bool +} + // A Connection is a logical bidirectional JSON-RPC connection. type Connection interface { // Read reads the next message to process off the connection. @@ -69,7 +102,7 @@ type Connection interface { type clientConnection interface { Connection - // SessionUpdated is called whenever the client session state changes. + // sessionUpdated is called whenever the client session state changes. sessionUpdated(clientSessionState) } @@ -81,16 +114,27 @@ type clientConnection interface { // TODO: should this interface be exported? type serverConnection interface { Connection + + // sessionUpdated is called whenever the server session state changes. sessionUpdated(ServerSessionState) } +// DefaultMaxLineLength is the default maximum number of bytes buffered while +// decoding a single inbound JSON-RPC frame. +const DefaultMaxLineLength = 16 * 1024 * 1024 + // A StdioTransport is a [Transport] that communicates over stdin/stdout using // newline-delimited JSON. -type StdioTransport struct{} +type StdioTransport struct { + // MaxLineLength bounds the number of bytes that may be buffered while + // decoding a single inbound JSON-RPC frame. A value of 0 selects [DefaultMaxLineLength], + // a negative value disables the cap. + MaxLineLength int +} // Connect implements the [Transport] interface. -func (*StdioTransport) Connect(context.Context) (Connection, error) { - return newIOConn(rwc{os.Stdin, nopCloserWriter{os.Stdout}}), nil +func (t *StdioTransport) Connect(context.Context) (Connection, error) { + return newIOConnLimited(rwc{os.Stdin, nopCloserWriter{os.Stdout}}, t.MaxLineLength), nil } // nopCloserWriter is an io.WriteCloser with a trivial Close method. @@ -105,11 +149,15 @@ func (nopCloserWriter) Close() error { return nil } type IOTransport struct { Reader io.ReadCloser Writer io.WriteCloser + // MaxLineLength bounds the number of bytes that may be buffered while + // decoding a single inbound JSON-RPC frame. A value of 0 selects [DefaultMaxLineLength], + // a negative value disables the cap. + MaxLineLength int } // Connect implements the [Transport] interface. func (t *IOTransport) Connect(context.Context) (Connection, error) { - return newIOConn(rwc{t.Reader, t.Writer}), nil + return newIOConnLimited(rwc{t.Reader, t.Writer}, t.MaxLineLength), nil } // An InMemoryTransport is a [Transport] that communicates over an in-memory @@ -147,7 +195,9 @@ type handler interface { handle(ctx context.Context, req *jsonrpc.Request) (any, error) } -func connect[H handler, State any](ctx context.Context, t Transport, b binder[H, State], s State, onClose func()) (H, error) { +// connect wires a transport to a binder. logger must be non-nil; it receives +// jsonrpc2 internal errors that would otherwise be dropped (see #218). +func connect[H handler, State any](ctx context.Context, t Transport, b binder[H, State], s State, onClose func(), logger *slog.Logger) (H, error) { var zero H mcpConn, err := t.Connect(ctx) if err != nil { @@ -157,13 +207,20 @@ func connect[H handler, State any](ctx context.Context, t Transport, b binder[H, reader, writer := jsonrpc2.Reader(mcpConn), jsonrpc2.Writer(mcpConn) var ( h H - preempter canceller + preempter = canceller{logger: logger} ) bind := func(conn *jsonrpc2.Connection) jsonrpc2.Handler { h = b.bind(mcpConn, conn, s, onClose) preempter.conn = conn return jsonrpc2.HandlerFunc(h.handle) } + // Transports may opt in to propagating cancellation of ctx into request + // handler contexts when their own lifecycle IS the cancellation signal + // (e.g., a connection bound to a single HTTP request). + var propagateCancellation bool + if cp, ok := mcpConn.(cancellationPropagator); ok { + propagateCancellation = cp.propagateCancellation() + } _ = jsonrpc2.NewConnection(ctx, jsonrpc2.ConnectionConfig{ Reader: reader, Writer: writer, @@ -173,57 +230,150 @@ func connect[H handler, State any](ctx context.Context, t Transport, b binder[H, OnDone: func() { b.disconnect(h) }, - OnInternalError: func(err error) { log.Printf("jsonrpc2 error: %v", err) }, + OnInternalError: func(err error) { + logger.Error("jsonrpc2 internal error", "error", err) + }, + PropagateCancellation: propagateCancellation, }) assert(preempter.conn != nil, "unbound preempter") return h, nil } +// cancellationPropagator is an optional interface implemented by a +// [Connection] whose own lifecycle should propagate cancellation into request +// handler contexts. The default jsonrpc2 behavior is to suppress propagation; +// transports that bind a connection to a single short-lived carrier (such as +// a one-shot HTTP request) should return true here so that handlers unwind +// when the carrier observes the peer going away. +type cancellationPropagator interface { + propagateCancellation() bool +} + // A canceller is a jsonrpc2.Preempter that cancels in-flight requests on MCP // cancelled notifications. type canceller struct { - conn *jsonrpc2.Connection + conn *jsonrpc2.Connection + logger *slog.Logger } // Preempt implements [jsonrpc2.Preempter]. func (c *canceller) Preempt(ctx context.Context, req *jsonrpc.Request) (result any, err error) { if req.Method == notificationCancelled { var params CancelledParams - if err := json.Unmarshal(req.Params, ¶ms); err != nil { + if err := internaljson.Unmarshal(req.Params, ¶ms); err != nil { return nil, err } id, err := jsonrpc2.MakeID(params.RequestID) if err != nil { return nil, err } - go c.conn.Cancel(id) + // The spec says implementations should log cancellation reasons, and + // the handler is the one place that can act on one, so the reason + // travels as the cause of the request's context rather than being + // dropped here. + c.logger.Debug("request cancelled by the peer", "id", id.Raw(), "reason", params.Reason) + go c.conn.CancelCause(id, &peerCancelledError{reason: params.Reason}) } return nil, jsonrpc2.ErrNotHandled } +// A peerCancelledError is the cause a request's context carries once the peer +// has sent a cancelled notification for it: [context.Cause] returns it to the +// handler with the reason the peer gave, and it unwraps to [context.Canceled] +// so that [errors.Is] keeps classifying the cancellation as one. +type peerCancelledError struct { + reason string +} + +func (e *peerCancelledError) Error() string { + if e.reason == "" { + return "request cancelled by the peer" + } + return "request cancelled by the peer: " + e.reason +} + +func (e *peerCancelledError) Unwrap() error { return context.Canceled } + +// callSubscriptionsListen issues a "subscriptions/listen" call (SEP-2575) +// without awaiting its JSON-RPC response. The call's logical lifetime is the +// stream of notifications that follow on the same channel — the empty +// response, if ever delivered, only marks subscription teardown — so the +// caller has nothing useful to block on. +// +// Cancellation is driven by ctx: when it is cancelled, a background goroutine +// sends a "notifications/cancelled" notification referencing the listen's +// request ID and retires the call from the connection's outgoing-calls map. +func callSubscriptionsListen(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params) { + call := conn.Call(ctx, method, params) + + go func() { + <-ctx.Done() + _ = cancelCall(ctx, conn, call) + }() +} + // call executes and awaits a jsonrpc2 call on the given connection, // translating errors into the mcp domain. func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params, result Result) error { - // TODO: the "%w"s in this function effectively make jsonrpc2.WireError part of the API. - // Consider alternatives. + // The "%w"s in this function expose jsonrpc.Error as part of the API. call := conn.Call(ctx, method, params) err := call.Await(ctx, result) switch { case errors.Is(err, jsonrpc2.ErrClientClosing), errors.Is(err, jsonrpc2.ErrServerClosing): return fmt.Errorf("%w: calling %q: %v", ErrConnectionClosed, method, err) case ctx.Err() != nil: - // Notify the peer of cancellation. - err := conn.Notify(xcontext.Detach(ctx), notificationCancelled, &CancelledParams{ - Reason: ctx.Err().Error(), - RequestID: call.ID().Raw(), - }) - return errors.Join(ctx.Err(), err) + // The notifications/cancelled message is best-effort. Retire the call + // immediately (so an unresponsive peer cannot delay the eager + // retirement that cancelCall's docstring promises) and send the + // notification off the caller's return path so a slow or unresponsive + // peer cannot delay the caller past its own deadline. See issue #1150. + // + // Setting MCPGODEBUG=blockingcancelnotify=1 restores the previous + // behavior of waiting synchronously for delivery inside cancelCall. + if blockingcancelnotify == "1" { + err := cancelCall(ctx, conn, call) + return errors.Join(ctx.Err(), err) + } + conn.Retire(call, ctx.Err()) + go func() { + notifyCtx, stop := context.WithTimeout(context.WithoutCancel(ctx), notifyCancellationTimeout) + defer stop() + _ = conn.Notify(notifyCtx, notificationCancelled, &CancelledParams{ + Reason: ctx.Err().Error(), + RequestID: call.ID().Raw(), + }) + }() + return ctx.Err() case err != nil: return fmt.Errorf("calling %q: %w", method, err) } return nil } +// cancelCall sends a "notifications/cancelled" notification for call and eagerly +// retires it from conn. +// +// By default, the jsonrpc2 library waits for graceful shutdown when the +// connection is closed, meaning it expects all outgoing and incoming requests +// to complete. However, for MCP this expectation is unrealistic, and can lead +// to hanging shutdown. For example, if a streamable client is killed, the +// server will not be able to detect this event, except via keepalive pings (if +// they are configured), and so outgoing calls may hang indefinitely. +// +// Therefore, we choose to eagerly retire calls, removing them from the +// outgoingCalls map, when the caller context is cancelled: if the caller will +// never receive the response, there's no need to track it. +func cancelCall(ctx context.Context, conn *jsonrpc2.Connection, call *jsonrpc2.AsyncCall) error { + notifyCtx, cancelNotify := context.WithTimeout(context.WithoutCancel(ctx), notifyCancellationTimeout) + defer cancelNotify() + err := conn.Notify(notifyCtx, notificationCancelled, &CancelledParams{ + Reason: ctx.Err().Error(), + RequestID: call.ID().Raw(), + }) + conn.Retire(call, ctx.Err()) + return err +} + // A LoggingTransport is a [Transport] that delegates to another transport, // writing RPC logs to an io.Writer. type LoggingTransport struct { @@ -328,7 +478,11 @@ func (r rwc) Close() error { // // See [msgBatch] for more discussion of message batching. type ioConn struct { - protocolVersion string // negotiated version, set during session initialization. + // protocolVersion is the negotiated version of the protocol, + // set during session initialization. + // Since writes may be concurrent to reads, we need to guard this with a mutex. + sessionMu sync.Mutex + protocolVersion string writeMu sync.Mutex // guards Write, which must be concurrency safe. rwc io.ReadWriteCloser // the underlying stream @@ -360,6 +514,17 @@ type msgOrErr struct { } func newIOConn(rwc io.ReadWriteCloser) *ioConn { + return newIOConnLimited(rwc, DefaultMaxLineLength) +} + +// newIOConnLimited builds an [ioConn] over rwc that bounds the number of bytes +// buffered while decoding a single inbound JSON-RPC frame to maxLineLength. +// maxLineLength == 0 selects [DefaultMaxLineLength], a negative value means no cap. +func newIOConnLimited(rwc io.ReadWriteCloser, maxLineLength int) *ioConn { + limit := maxLineLength + if limit == 0 { + limit = DefaultMaxLineLength + } var ( incoming = make(chan msgOrErr) closed = make(chan struct{}) @@ -371,7 +536,15 @@ func newIOConn(rwc io.ReadWriteCloser) *ioConn { // but that is unavoidable since AFAIK there is no (easy and portable) way to // guarantee that reads of stdin are unblocked when closed. go func() { - dec := json.NewDecoder(rwc) + var ( + reader io.Reader = rwc + limiter *frameLimitReader + ) + if limit > 0 { + limiter = &frameLimitReader{r: rwc, limit: limit} + reader = limiter + } + dec := json.NewDecoder(reader) for { var raw json.RawMessage err := dec.Decode(&raw) @@ -381,7 +554,8 @@ func newIOConn(rwc io.ReadWriteCloser) *ioConn { var tr [1]byte if n, readErr := dec.Buffered().Read(tr[:]); n > 0 { // If read byte is not a newline, it is an error. - if tr[0] != '\n' { + // Support both Unix (\n) and Windows (\r\n) line endings. + if tr[0] != '\n' && tr[0] != '\r' { err = fmt.Errorf("invalid trailing data at the end of stream") } } else if readErr != nil && readErr != io.EOF { @@ -396,6 +570,9 @@ func newIOConn(rwc io.ReadWriteCloser) *ioConn { if err != nil { return } + if limiter != nil { + limiter.resetFrame() + } } }() return &ioConn{ @@ -405,17 +582,45 @@ func newIOConn(rwc io.ReadWriteCloser) *ioConn { } } +// errFrameTooLarge means that a single inbound JSON-RPC frame exceeded the configured byte +// limit before the value was completely received. +var errFrameTooLarge = errors.New("inbound JSON-RPC frame exceeded the configured maximum line length") + +// frameLimitReader bounds the number of bytes [json.Decoder] may buffer while +// decoding a single JSON value. Read returns [errFrameTooLarge] once the budget is exhausted. +type frameLimitReader struct { + r io.Reader + limit int + count int +} + +func (r *frameLimitReader) Read(p []byte) (int, error) { + if r.count >= r.limit { + return 0, errFrameTooLarge + } + if len(p) > r.limit-r.count { + p = p[:r.limit-r.count] + } + n, err := r.r.Read(p) + r.count += n + return n, err +} + +func (r *frameLimitReader) resetFrame() { r.count = 0 } + func (c *ioConn) SessionID() string { return "" } func (c *ioConn) sessionUpdated(state ServerSessionState) { - protocolVersion := "" - if state.InitializeParams != nil { - protocolVersion = state.InitializeParams.ProtocolVersion - } + protocolVersion := state.NegotiatedProtocolVersion if protocolVersion == "" { + // 2025-03-26 is used, because it's the last spec version + // where specifying the protocol version in the HTTP header + // was not required. protocolVersion = protocolVersion20250326 } - c.protocolVersion = negotiatedVersion(protocolVersion) + c.sessionMu.Lock() + c.protocolVersion = protocolVersion + c.sessionMu.Unlock() } // addBatch records a msgBatch for an incoming batch payload. @@ -516,8 +721,12 @@ func (t *ioConn) Read(ctx context.Context) (jsonrpc.Message, error) { if err != nil { return nil, err } - if batch && t.protocolVersion >= protocolVersion20250618 { - return nil, fmt.Errorf("JSON-RPC batching is not supported in %s and later (request version: %s)", protocolVersion20250618, t.protocolVersion) + var protocolVersion string + t.sessionMu.Lock() + protocolVersion = t.protocolVersion + t.sessionMu.Unlock() + if batch && protocolVersion >= protocolVersion20250618 { + return nil, fmt.Errorf("JSON-RPC batching is not supported in %s and later (request version: %s)", protocolVersion20250618, protocolVersion) } t.queue = msgs[1:] @@ -525,7 +734,7 @@ func (t *ioConn) Read(ctx context.Context) (jsonrpc.Message, error) { if batch { var respBatch *msgBatch // track incoming requests in the batch for _, msg := range msgs { - if req, ok := msg.(*jsonrpc.Request); ok { + if req, ok := msg.(*jsonrpc.Request); ok && req.IsCall() { if respBatch == nil { respBatch = &msgBatch{ unresolved: make(map[jsonrpc2.ID]int), @@ -553,7 +762,7 @@ func (t *ioConn) Read(ctx context.Context) (jsonrpc.Message, error) { func readBatch(data []byte) (msgs []jsonrpc.Message, isBatch bool, _ error) { // Try to read an array of messages first. var rawBatch []json.RawMessage - if err := json.Unmarshal(data, &rawBatch); err == nil { + if err := internaljson.Unmarshal(data, &rawBatch); err == nil { if len(rawBatch) == 0 { return nil, true, fmt.Errorf("empty batch") } diff --git a/mcp/transport_example_test.go b/mcp/transport_example_test.go index 7390ea4..e264133 100644 --- a/mcp/transport_example_test.go +++ b/mcp/transport_example_test.go @@ -2,9 +2,6 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. -// Uses strings.SplitSeq. -//go:build go1.24 - package mcp_test import ( @@ -45,9 +42,8 @@ func ExampleLoggingTransport() { } // Output: - // read: {"jsonrpc":"2.0","id":1,"result":{"capabilities":{"logging":{}},"protocolVersion":"2025-06-18","serverInfo":{"name":"server","version":"v0.0.1"}}} - // write: {"jsonrpc":"2.0","id":1,"method":"initialize","params":{"capabilities":{"roots":{"listChanged":true}},"clientInfo":{"name":"client","version":"v0.0.1"},"protocolVersion":"2025-06-18"}} - // write: {"jsonrpc":"2.0","method":"notifications/initialized","params":{}} + // read: {"jsonrpc":"2.0","id":1,"result":{"resultType":"complete","_meta":{"io.modelcontextprotocol/serverInfo":{"name":"server","version":"v0.0.1"}},"ttlMs":0,"cacheScope":"public","supportedVersions":["2026-07-28","2025-11-25","2025-06-18","2025-03-26","2024-11-05"],"capabilities":{"logging":{}}}} + // write: {"jsonrpc":"2.0","id":1,"method":"server/discover","params":{"_meta":{"io.modelcontextprotocol/clientCapabilities":{"roots":{"listChanged":true}},"io.modelcontextprotocol/clientInfo":{"name":"client","version":"v0.0.1"},"io.modelcontextprotocol/protocolVersion":"2026-07-28"}}} } // !-loggingtransport diff --git a/mcp/transport_test.go b/mcp/transport_test.go index aeff366..a20ee5f 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -5,10 +5,14 @@ package mcp import ( + "bytes" "context" + "errors" + "fmt" "io" "strings" "testing" + "time" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" "github.com/modelcontextprotocol/go-sdk/jsonrpc" @@ -56,10 +60,15 @@ func TestBatchFraming(t *testing.T) { func TestIOConnRead(t *testing.T) { tests := []struct { - name string - input string - want string + name string + input string + want string + // protocolVersion is the version negotiated by 'initialize'; empty + // means no session state was pushed to the connection. protocolVersion string + // requested is the version the client asked for, which differs from + // protocolVersion when the server counter-offered. + requested string }{ { name: "valid json input", @@ -83,6 +92,11 @@ func TestIOConnRead(t *testing.T) { want: "", protocolVersion: "", }, + { + name: "windows newline at the end of first valid json input", + input: "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"test\",\"params\":{}}\r\n", + want: "", + }, { name: "batching old protocol", input: `[{"jsonrpc":"2.0","id":1,"method":"test1"},{"jsonrpc":"2.0","id":2,"method":"test2"}]`, @@ -95,6 +109,16 @@ func TestIOConnRead(t *testing.T) { want: "JSON-RPC batching is not supported in 2025-06-18 and later (request version: 2025-06-18)", protocolVersion: protocolVersion20250618, }, + { + // The client asked for a version the server does not negotiate, so + // the connection must follow the server's counter-offer (which + // forbids batching) rather than the client's request. + name: "batching at a counter-offered version", + input: `[{"jsonrpc":"2.0","id":1,"method":"test1"},{"jsonrpc":"2.0","id":2,"method":"test2"}]`, + want: "JSON-RPC batching is not supported in 2025-06-18 and later (request version: 2025-11-25)", + requested: protocolVersion20241105, + protocolVersion: protocolVersion20251125, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -103,10 +127,15 @@ func TestIOConnRead(t *testing.T) { }) t.Cleanup(func() { tr.Close() }) if tt.protocolVersion != "" { + requested := tt.requested + if requested == "" { + requested = tt.protocolVersion + } tr.sessionUpdated(ServerSessionState{ InitializeParams: &InitializeParams{ - ProtocolVersion: tt.protocolVersion, + ProtocolVersion: requested, }, + NegotiatedProtocolVersion: tt.protocolVersion, }) } _, err := tr.Read(context.Background()) @@ -119,3 +148,210 @@ func TestIOConnRead(t *testing.T) { }) } } + +// go-sdk#976: stdio must surface empty-method calls as requests, not responses. +func TestIOConnRead_EmptyMethod(t *testing.T) { + tr := newIOConn(rwc{ + rc: io.NopCloser(strings.NewReader(`{"jsonrpc":"2.0","id":5,"method":"","params":{}}`)), + }) + t.Cleanup(func() { tr.Close() }) + + msg, err := tr.Read(context.Background()) + if err != nil { + t.Fatalf("ioConn.Read() error = %v", err) + } + req, ok := msg.(*jsonrpc.Request) + if !ok { + t.Fatalf("message type = %T, want *jsonrpc.Request", msg) + } + if req.Method != "" { + t.Errorf("Method = %q, want empty string", req.Method) + } + if req.ID != jsonrpc2.Int64ID(5) { + t.Errorf("ID = %v, want 5", req.ID.Raw()) + } +} + +// TestIOConnReadBatchNotifications is a regression test for go-sdk#1256: +// notifications carried inside a JSON-RPC batch must not be tracked as awaiting +// a response. Only calls belong in a batch's response array, so a batch's +// notifications must neither keep it from completing nor pollute the +// per-connection duplicate-ID tracking. +func TestIOConnReadBatchNotifications(t *testing.T) { + ctx := context.Background() + + tests := []struct { + name string + // input is one or more newline-separated batch payloads. + input string + // wantCalls records, for each message Read returns in order, whether it + // is a call (true) or a notification (false). Reading every message + // must succeed. + wantCalls []bool + // wantTrackedCalls is the number of in-flight call requests tracked across + // batches awaiting a response once every message has been read, before any + // response is written. + wantTrackedCalls int + // respondIDs are the call IDs to respond to, in order. + respondIDs []int64 + // wantWrites are the batch response payloads expected on the wire, each + // matched as a substring. + wantWrites []string + }{ + { + name: "call and notification is answered", + input: `[{"jsonrpc":"2.0","id":2,"method":"ping"},{"jsonrpc":"2.0","method":"notifications/cancelled","params":{"requestId":99}}]`, + wantCalls: []bool{true, false}, + wantTrackedCalls: 1, + respondIDs: []int64{2}, + wantWrites: []string{`[{"jsonrpc":"2.0","id":2,"result":{}}]`}, + }, + { + name: "only notifications keeps the connection open", + input: `[{"jsonrpc":"2.0","method":"notifications/cancelled","params":{"requestId":98}},{"jsonrpc":"2.0","method":"notifications/cancelled","params":{"requestId":99}}]`, + wantCalls: []bool{false, false}, + wantTrackedCalls: 0, + }, + { + name: "repeated batches holding a notification", + input: `[{"jsonrpc":"2.0","id":2,"method":"ping"},{"jsonrpc":"2.0","method":"notifications/cancelled","params":{"requestId":99}}]` + "\n" + + `[{"jsonrpc":"2.0","id":3,"method":"ping"},{"jsonrpc":"2.0","method":"notifications/cancelled","params":{"requestId":99}}]`, + wantCalls: []bool{true, false, true, false}, + wantTrackedCalls: 2, + respondIDs: []int64{2, 3}, + wantWrites: []string{ + `[{"jsonrpc":"2.0","id":2,"result":{}}]`, + `[{"jsonrpc":"2.0","id":3,"result":{}}]`, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sink := new(bytes.Buffer) + tr := newIOConn(rwc{ + rc: io.NopCloser(strings.NewReader(tt.input)), + wc: nopCloserWriter{sink}, + }) + t.Cleanup(func() { tr.Close() }) + + for i, wantCall := range tt.wantCalls { + msg, err := tr.Read(ctx) + if err != nil { + t.Fatalf("Read() #%d: unexpected error: %v", i, err) + } + req, ok := msg.(*jsonrpc.Request) + if !ok { + t.Fatalf("Read() #%d = %T, want *jsonrpc.Request", i, msg) + } + if req.IsCall() != wantCall { + t.Fatalf("Read() #%d IsCall() = %t, want %t", i, req.IsCall(), wantCall) + } + } + + // Only calls are tracked as awaiting a response; notifications must + // not linger in the batch bookkeeping. + tr.batchMu.Lock() + gotTrackedCalls := len(tr.batches) + tr.batchMu.Unlock() + if gotTrackedCalls != tt.wantTrackedCalls { + t.Errorf("tracked calls = %d, want %d", gotTrackedCalls, tt.wantTrackedCalls) + } + + for _, id := range tt.respondIDs { + resp := &jsonrpc.Response{ID: jsonrpc2.Int64ID(id), Result: []byte("{}")} + if err := tr.Write(ctx, resp); err != nil { + t.Fatalf("Write(response id=%d): %v", id, err) + } + } + + got := sink.String() + for _, want := range tt.wantWrites { + if !strings.Contains(got, want) { + t.Errorf("batch responses = %q, missing %q", got, want) + } + } + }) + } +} + +// TestIOConnFrameCap verifies single inbound frame cap enforcement. +func TestIOConnFrameCap(t *testing.T) { + tests := []struct { + name string + r io.ReadCloser + limit int + wantMessageCount int + wantErr bool + }{ + { + name: "infinite string", + r: &endlessReader{prefix: `"`, repeat: "A"}, + limit: 1024, + wantErr: true, + }, + { + name: "cap does not reset per new line", + r: &endlessReader{prefix: "[", repeat: "0,\n"}, + limit: 1024, + wantErr: true, + }, + { + name: "cap resets per message", + limit: 1024, // < 600 * 3 + r: io.NopCloser(strings.NewReader(func() string { + var b strings.Builder + for i := range 3 { + fmt.Fprintf(&b, `{"jsonrpc":"2.0","id":%d,"method":"test","params":{"pad":"`, i) + b.WriteString(strings.Repeat("A", 600)) + b.WriteString(`"}}`) + b.WriteByte('\n') + } + return b.String() + }())), + wantMessageCount: 3, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tr := newIOConnLimited(rwc{rc: tt.r}, tt.limit) + t.Cleanup(func() { tr.Close() }) + + read := func() (jsonrpc.Message, error) { + type result struct { + msg jsonrpc.Message + err error + } + done := make(chan result, 1) + + go func() { + msg, err := tr.Read(context.Background()) + done <- result{msg, err} + }() + + select { + case r := <-done: + return r.msg, r.err + case <-time.After(10 * time.Second): + t.Fatal("Read() did not return: frame cap failed to trip") + return nil, nil + } + } + + for i := range tt.wantMessageCount { + msg, err := read() + if err != nil { + t.Fatalf("Read() #%d error = %v, want nil", i, err) + } + if msg == nil { + t.Fatalf("Read() #%d returned nil message", i) + } + } + if tt.wantErr { + if _, err := read(); !errors.Is(err, errFrameTooLarge) { + t.Fatalf("Read() error = %v, want errFrameTooLarge", err) + } + } + }) + } +} diff --git a/mcp/util.go b/mcp/util.go index 5ada466..8ffaa74 100644 --- a/mcp/util.go +++ b/mcp/util.go @@ -5,8 +5,9 @@ package mcp import ( - "crypto/rand" "encoding/json" + + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" ) func assert(cond bool, msg string) { @@ -15,20 +16,6 @@ func assert(cond bool, msg string) { } } -// Copied from crypto/rand. -// TODO: once 1.24 is assured, just use crypto/rand. -const base32alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567" - -func randText() string { - // ⌈log₃₂ 2¹²⁸⌉ = 26 chars - src := make([]byte, 26) - rand.Read(src) - for i := range src { - src[i] = base32alphabet[src[i]%32] - } - return string(src) -} - // remarshal marshals from to JSON, and then unmarshals into to, which must be // a pointer type. func remarshal(from, to any) error { @@ -36,7 +23,7 @@ func remarshal(from, to any) error { if err != nil { return err } - if err := json.Unmarshal(data, to); err != nil { + if err := internaljson.Unmarshal(data, to); err != nil { return err } return nil diff --git a/mkdocs.yaml b/mkdocs.yaml new file mode 100644 index 0000000..e978c85 --- /dev/null +++ b/mkdocs.yaml @@ -0,0 +1,98 @@ +site_name: MCP Go SDK +site_url: https://go.sdk.modelcontextprotocol.io/ +site_description: Golang SDK for the Model Context Protocol - standardized integration between AI models and tools + +repo_url: https://github.com/modelcontextprotocol/go-sdk +repo_name: modelcontextprotocol/go-sdk +edit_uri: edit/main/docs/ + +theme: + name: material + favicon: images/logo-dark.svg + logo: images/logo-light.svg + palette: + - media: "(prefers-color-scheme: dark)" + scheme: slate + primary: cyan + accent: teal + toggle: + icon: material/lightbulb + name: "Switch to light mode" + - media: "(prefers-color-scheme: light)" + scheme: default + primary: cyan + accent: teal + toggle: + icon: material/lightbulb-outline + name: "Switch to dark mode" + + features: + - navigation.instant + - navigation.instant.progress + - navigation.tabs + - navigation.tabs.sticky + - navigation.sections + - navigation.top + - navigation.path + - navigation.indexes + - toc.follow + - search.suggest + - search.highlight + - content.code.copy + - content.code.annotate + + +nav: + - Documentation: + - Overview: README.md + - Quick Start: quick_start.md + - LifeCycle: protocol.md + - MCP Components: + - MCP Client: client.md + - MCP Server: server.md + - Troubleshooting: troubleshooting.md + - Backwards Compatibility: mcpgodebug.md + - Rough Edges: rough_edges.md + - Contributing: + - Contributing Guide: contributing.md + - API Reference: https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk + + +markdown_extensions: + - admonition + - pymdownx.details + - pymdownx.superfences: + custom_fences: + - name: mermaid + class: mermaid + format: !!python/name:pymdownx.superfences.fence_code_format + - pymdownx.tabbed: + alternate_style: true + - pymdownx.highlight: + anchor_linenums: true + line_spans: __span + - pymdownx.inlinehilite + - pymdownx.snippets: + base_path: . + - pymdownx.mark + - pymdownx.critic + - pymdownx.caret + - pymdownx.keys + - pymdownx.tilde + - pymdownx.emoji: + emoji_index: !!python/name:material.extensions.emoji.twemoji + emoji_generator: !!python/name:material.extensions.emoji.to_svg + - attr_list + - md_in_html + - tables + - toc: + permalink: true + +extra: + social: + - icon: fontawesome/brands/github + link: https://github.com/modelcontextprotocol/go-sdk + generator: false + +plugins: + - search \ No newline at end of file diff --git a/oauthex/audience.go b/oauthex/audience.go new file mode 100644 index 0000000..cad1303 --- /dev/null +++ b/oauthex/audience.go @@ -0,0 +1,44 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package oauthex + +import "strings" + +// MatchesResource reports whether any of claims matches resource under +// RFC 3986 §6.2.3 scheme-based normalization, narrowed to the empty-path +// case: a URI with an empty path is treated as equivalent to one with a +// path of "/". All other URI components (scheme, host case, port, query, +// fragment) must match exactly — this is intentionally narrower than the +// full §6.2.3 rung, which would also fold scheme/host case and default +// ports. +// +// RFC 9728 §3.3 normatively requires "simple string comparison" per +// RFC 3986 §6.2.1 (byte-equal). Callers that need strict byte-equal +// semantics should compare claims to resource directly: +// +// for _, c := range claims { if c == resource { return true } } +// +// Background: RFC 9728 §3.3 canonicalises the protected-resource +// identifier with a trailing slash, but RFC 8707 resource indicators +// sometimes omit it, and upstream IdPs vary in which form they emit in +// `aud` claims (Google trims, Auth0 retains, claude.ai round-trips +// whichever it received). Strict byte equality therefore fails routinely +// on legitimate setups; this helper is the most common pragmatic relaxation +// while keeping path/scheme/host strict so token confusion across distinct +// resources still fails closed. +// +// Returns false when claims is empty. +func MatchesResource(claims []string, resource string) bool { + expected := strings.TrimSuffix(resource, "/") + for _, c := range claims { + if c == resource { + return true + } + if strings.TrimSuffix(c, "/") == expected { + return true + } + } + return false +} diff --git a/oauthex/audience_test.go b/oauthex/audience_test.go new file mode 100644 index 0000000..7d5acfb --- /dev/null +++ b/oauthex/audience_test.go @@ -0,0 +1,101 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package oauthex + +import "testing" + +func TestMatchesResource(t *testing.T) { + tests := []struct { + name string + claims []string + resource string + want bool + }{ + { + name: "exact match", + claims: []string{"https://mcp.example.com/"}, + resource: "https://mcp.example.com/", + want: true, + }, + { + name: "claim has trailing slash, resource does not", + claims: []string{"https://mcp.example.com/"}, + resource: "https://mcp.example.com", + want: true, + }, + { + name: "claim missing trailing slash, resource has one", + claims: []string{"https://mcp.example.com"}, + resource: "https://mcp.example.com/", + want: true, + }, + { + name: "multiple claims, one matches", + claims: []string{"https://other.example.com/", "https://mcp.example.com"}, + resource: "https://mcp.example.com/", + want: true, + }, + { + name: "no claims", + claims: nil, + resource: "https://mcp.example.com/", + want: false, + }, + { + name: "claim does not match resource", + claims: []string{"https://attacker.example.com/"}, + resource: "https://mcp.example.com/", + want: false, + }, + { + name: "path mismatch is not tolerated", + claims: []string{"https://mcp.example.com/v2"}, + resource: "https://mcp.example.com", + want: false, + }, + { + name: "scheme mismatch is not tolerated", + claims: []string{"http://mcp.example.com/"}, + resource: "https://mcp.example.com/", + want: false, + }, + // The following document the intentional boundaries: §6.2.3 also + // permits scheme/host case folding and default-port elision, but + // MatchesResource deliberately does NOT, so two distinct + // registered resources cannot collide via these normalizations. + { + name: "host case difference is not tolerated", + claims: []string{"https://MCP.example.com/"}, + resource: "https://mcp.example.com/", + want: false, + }, + { + name: "default-port elision is not tolerated", + claims: []string{"https://mcp.example.com:443/"}, + resource: "https://mcp.example.com/", + want: false, + }, + { + name: "query string difference is not tolerated", + claims: []string{"https://mcp.example.com/?x=y"}, + resource: "https://mcp.example.com/", + want: false, + }, + { + name: "surrounding whitespace is not tolerated (malformed claim)", + claims: []string{" https://mcp.example.com/ "}, + resource: "https://mcp.example.com/", + want: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := MatchesResource(tt.claims, tt.resource) + if got != tt.want { + t.Errorf("MatchesResource(%v, %q) = %v, want %v", tt.claims, tt.resource, got, tt.want) + } + }) + } +} diff --git a/oauthex/auth_meta.go b/oauthex/auth_meta.go index 3723318..711b17a 100644 --- a/oauthex/auth_meta.go +++ b/oauthex/auth_meta.go @@ -5,8 +5,6 @@ // This file implements Authorization Server Metadata. // See https://www.rfc-editor.org/rfc/rfc8414.html. -//go:build mcp_go_client_oauth - package oauthex import ( @@ -14,6 +12,8 @@ import ( "errors" "fmt" "net/http" + + "github.com/modelcontextprotocol/go-sdk/internal/authutil" ) // AuthServerMeta represents the metadata for an OAuth 2.0 authorization server, @@ -22,10 +22,12 @@ import ( // Not supported: // - signed metadata // +// Note: URL fields in this struct are validated by validateAuthServerMetaURLs to +// prevent XSS attacks. If you add a new URL field, you must also add it to that +// function. +// // [RFC 8414]: https://tools.ietf.org/html/rfc8414) type AuthServerMeta struct { - // GENERATED BY GEMINI 2.5. - // Issuer is the REQUIRED URL identifying the authorization server. Issuer string `json:"issuer"` @@ -109,44 +111,106 @@ type AuthServerMeta struct { // CodeChallengeMethodsSupported is a RECOMMENDED JSON array of strings containing a list of // PKCE code challenge methods supported by this authorization server. CodeChallengeMethodsSupported []string `json:"code_challenge_methods_supported,omitempty"` -} -var wellKnownPaths = []string{ - "/.well-known/oauth-authorization-server", - "/.well-known/openid-configuration", + // ClientIDMetadataDocumentSupported is a boolean indicating whether the authorization server + // supports client ID metadata documents. + ClientIDMetadataDocumentSupported bool `json:"client_id_metadata_document_supported,omitempty"` + + // AuthorizationResponseIssParameterSupported indicates whether the authorization server + // provides the "iss" parameter in authorization responses per [RFC 9207]. + // When true, clients must verify the "iss" parameter is present and matches the Issuer field. + // + // [RFC 9207]: https://www.rfc-editor.org/rfc/rfc9207 + AuthorizationResponseIssParameterSupported bool `json:"authorization_response_iss_parameter_supported,omitempty"` } // GetAuthServerMeta issues a GET request to retrieve authorization server metadata -// from an OAuth authorization server with the given issuerURL. +// from an OAuth authorization server with the given metadataURL. // // It follows [RFC 8414]: -// - The well-known paths specified there are inserted into the URL's path, one at time. -// The first to succeed is used. -// - The Issuer field is checked against issuerURL. +// - The metadataURL must use HTTPS or be a local address. +// - The Issuer field is checked against metadataURL.Issuer. +// +// It also verifies that the authorization server supports PKCE and that the URLs +// in the metadata don't use dangerous schemes. +// +// It returns an error if the request fails with a non-4xx status code or the fetched +// metadata doesn't pass security validations. +// It returns nil if the request fails with a 4xx status code. // // [RFC 8414]: https://tools.ietf.org/html/rfc8414 -func GetAuthServerMeta(ctx context.Context, issuerURL string, c *http.Client) (*AuthServerMeta, error) { - var errs []error - for _, p := range wellKnownPaths { - u, err := prependToPath(issuerURL, p) - if err != nil { - // issuerURL is bad; no point in continuing. - return nil, err - } - asm, err := getJSON[AuthServerMeta](ctx, c, u, 1<<20) - if err == nil { - if asm.Issuer != issuerURL { // section 3.3 - // Security violation; don't keep trying. - return nil, fmt.Errorf("metadata issuer %q does not match issuer URL %q", asm.Issuer, issuerURL) - } - - if len(asm.CodeChallengeMethodsSupported) == 0 { - return nil, fmt.Errorf("authorization server at %s does not implement PKCE", issuerURL) - } - - return asm, nil - } - errs = append(errs, err) +func GetAuthServerMeta(ctx context.Context, metadataURL, issuer string, c *http.Client) (*AuthServerMeta, error) { + // Only allow HTTP for local addresses (testing or development purposes). + if err := checkHTTPSOrLoopback(metadataURL); err != nil { + return nil, fmt.Errorf("metadataURL: %v", err) } - return nil, fmt.Errorf("failed to get auth server metadata from %q: %w", issuerURL, errors.Join(errs...)) + asm, err := getJSON[AuthServerMeta](ctx, c, metadataURL, 1<<20) + if err != nil { + var httpErr *httpStatusError + if errors.As(err, &httpErr) { + if 400 <= httpErr.StatusCode && httpErr.StatusCode < 500 { + return nil, nil + } + } + return nil, fmt.Errorf("%v", err) // Do not expose error types. + } + if !authutil.IssuersEqual(asm.Issuer, issuer) { + return nil, fmt.Errorf("metadata issuer %q does not match issuer URL %q", asm.Issuer, issuer) + } + + if len(asm.CodeChallengeMethodsSupported) == 0 { + return nil, fmt.Errorf("authorization server at %s does not implement PKCE", issuer) + } + + // Validate endpoint URLs to prevent XSS attacks (see #526). + if err := validateAuthServerMetaURLs(asm); err != nil { + return nil, err + } + + return asm, nil +} + +// validateAuthServerMetaURLs validates all URL fields in AuthServerMeta +// to ensure they don't use dangerous schemes that could enable XSS attacks. +// It also validates that URLs likely to be called by the client use +// HTTPS or are loopback addresses. +func validateAuthServerMetaURLs(asm *AuthServerMeta) error { + urls := []struct { + name string + value string + }{ + {"authorization_endpoint", asm.AuthorizationEndpoint}, + {"token_endpoint", asm.TokenEndpoint}, + {"jwks_uri", asm.JWKSURI}, + {"registration_endpoint", asm.RegistrationEndpoint}, + {"service_documentation", asm.ServiceDocumentation}, + {"op_policy_uri", asm.OpPolicyURI}, + {"op_tos_uri", asm.OpTOSURI}, + {"revocation_endpoint", asm.RevocationEndpoint}, + {"introspection_endpoint", asm.IntrospectionEndpoint}, + } + + for _, u := range urls { + if err := checkURLScheme(u.value); err != nil { + return fmt.Errorf("%s: %w", u.name, err) + } + } + + urls = []struct { + name string + value string + }{ + {"authorization_endpoint", asm.AuthorizationEndpoint}, + {"token_endpoint", asm.TokenEndpoint}, + {"registration_endpoint", asm.RegistrationEndpoint}, + {"introspection_endpoint", asm.IntrospectionEndpoint}, + } + + for _, u := range urls { + if err := checkHTTPSOrLoopback(u.value); err != nil { + return fmt.Errorf("%s: %w", u.name, err) + } + } + + return nil } diff --git a/oauthex/auth_meta_test.go b/oauthex/auth_meta_test.go index 1e60882..5c6c4ad 100644 --- a/oauthex/auth_meta_test.go +++ b/oauthex/auth_meta_test.go @@ -2,8 +2,6 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. -//go:build mcp_go_client_oauth - package oauthex import ( @@ -37,9 +35,10 @@ func TestAuthMetaParse(t *testing.T) { func TestGetAuthServerMetaPKCESupport(t *testing.T) { ctx := context.Background() tests := []struct { - name string - hasPKCESupport bool - wantError string + name string + hasPKCESupport bool + wantError string + issuerWithTrailingSlash bool }{ { name: "server_with_pkce_support", @@ -50,6 +49,12 @@ func TestGetAuthServerMetaPKCESupport(t *testing.T) { hasPKCESupport: false, wantError: "does not implement PKCE", }, + { + // ProtectedResourceMetadata may contain AuthorizationServers with a trailing slash (see Issue #953) + name: "issuer_with_trailing_slash", + hasPKCESupport: true, + issuerWithTrailingSlash: true, + }, } for _, tt := range tests { @@ -86,6 +91,10 @@ func TestGetAuthServerMetaPKCESupport(t *testing.T) { // The fake server sets issuer to https://localhost:, so compute that issuer. u, _ := url.Parse(ts.URL) issuer := "https://localhost:" + u.Port() + metadataURL := issuer + "/.well-known/oauth-authorization-server" + if tt.issuerWithTrailingSlash { + issuer += "/" + } // The fake server presents a cert for example.com; set ServerName accordingly. httpClient := ts.Client() @@ -95,7 +104,7 @@ func TestGetAuthServerMetaPKCESupport(t *testing.T) { httpClient.Transport = clone } - meta, err := GetAuthServerMeta(ctx, issuer, httpClient) + meta, err := GetAuthServerMeta(ctx, metadataURL, issuer, httpClient) if tt.wantError != "" { if err == nil { t.Fatal("wanted error but got none") diff --git a/oauthex/client.go b/oauthex/client.go new file mode 100644 index 0000000..1b19eed --- /dev/null +++ b/oauthex/client.go @@ -0,0 +1,66 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package oauthex + +import "errors" + +// ClientCredentials holds client authentication credentials for OAuth token requests. +// It supports multiple authentication methods, but only one method should be set at a time. +// Use the Validate method to ensure proper configuration. +type ClientCredentials struct { + // ClientID is the OAuth2 client identifier. + // REQUIRED for all authentication methods. + ClientID string + + // ClientSecretAuth configures client authentication using a client secret. + // This is the most common authentication method for confidential clients. + // OPTIONAL. If not provided, the client is treated as a public client. + ClientSecretAuth *ClientSecretAuth + + // Issuer is the issuer identifier of the authorization server these + // credentials are registered with. Pre-registered credentials are bound + // to a specific authorization server; when set, an error is returned if + // the discovered authorization server does not match, per SEP-2352. + // The comparison ignores a single trailing slash, matching the + // tolerance applied during RFC 8414 Section 3.3 metadata validation. + // OPTIONAL. + Issuer string +} + +// ClientSecretAuth holds client secret authentication credentials. +// This authentication method supports both "client_secret_basic" and "client_secret_post" +// methods as defined in RFC 6749 Section 2.3.1. +type ClientSecretAuth struct { + // ClientSecret is the OAuth2 client secret for confidential clients. + // REQUIRED when using ClientSecretAuth. + ClientSecret string +} + +// Validate checks that the ClientCredentials are properly configured. +// It ensures that: +// - ClientID is not empty. +// - At most one authentication method is configured. +// - If ClientSecretAuth is set, ClientSecret is not empty. +func (c *ClientCredentials) Validate() error { + if c.ClientID == "" { + return errors.New("ClientID is required") + } + + // Count how many auth methods are configured. + authMethodCount := 0 + if c.ClientSecretAuth != nil { + authMethodCount++ + if c.ClientSecretAuth.ClientSecret == "" { + return errors.New("ClientSecret is required when using ClientSecretAuth") + } + } + + // Allow zero auth methods (public client) or exactly one auth method. + if authMethodCount > 1 { + return errors.New("only one client authentication method can be configured") + } + + return nil +} diff --git a/oauthex/client_test.go b/oauthex/client_test.go new file mode 100644 index 0000000..34d8188 --- /dev/null +++ b/oauthex/client_test.go @@ -0,0 +1,88 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package oauthex + +import ( + "reflect" + "strings" + "testing" +) + +func TestClientCredentials_Validate(t *testing.T) { + tests := []struct { + name string + creds ClientCredentials + wantErr string + }{ + { + name: "valid public client", + creds: ClientCredentials{ClientID: "my-client"}, + }, + { + name: "valid confidential client", + creds: ClientCredentials{ + ClientID: "my-client", + ClientSecretAuth: &ClientSecretAuth{ + ClientSecret: "my-secret", + }, + }, + }, + { + name: "empty client ID", + creds: ClientCredentials{}, + wantErr: "ClientID is required", + }, + { + name: "empty secret in ClientSecretAuth", + creds: ClientCredentials{ + ClientID: "my-client", + ClientSecretAuth: &ClientSecretAuth{}, + }, + wantErr: "ClientSecret is required", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.creds.Validate() + if tt.wantErr == "" { + if err != nil { + t.Fatalf("Validate() unexpected error: %v", err) + } + return + } + if err == nil { + t.Fatalf("Validate() expected error containing %q, got nil", tt.wantErr) + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("Validate() error = %v, want error containing %q", err, tt.wantErr) + } + }) + } +} + +// TestClientCredentials_ValidateCoversAllAuthFields uses reflection to detect +// when new authentication method fields are added to ClientCredentials without +// updating Validate. If this test fails, update Validate() to handle the new +// field and increment knownAuthMethods. +func TestClientCredentials_ValidateCoversAllAuthFields(t *testing.T) { + const knownAuthMethods = 1 // ClientSecretAuth + + typ := reflect.TypeOf(ClientCredentials{}) + var pointerFields int + for i := range typ.NumField() { + f := typ.Field(i) + if f.Name == "ClientID" || f.Name == "Issuer" { + continue + } + if f.Type.Kind() != reflect.Ptr { + t.Errorf("field %q is %v, expected a pointer to an auth method struct", f.Name, f.Type.Kind()) + } + pointerFields++ + } + + if pointerFields != knownAuthMethods { + t.Fatalf("ClientCredentials has %d auth method fields but Validate only knows about %d -- update Validate() and this test", pointerFields, knownAuthMethods) + } +} diff --git a/oauthex/dcr.go b/oauthex/dcr.go index 75ce296..f95c55b 100644 --- a/oauthex/dcr.go +++ b/oauthex/dcr.go @@ -5,8 +5,6 @@ // This file implements Authorization Server Metadata. // See https://www.rfc-editor.org/rfc/rfc8414.html. -//go:build mcp_go_client_oauth - package oauthex import ( @@ -17,9 +15,15 @@ import ( "io" "net/http" "time" + + internaljson "github.com/modelcontextprotocol/go-sdk/internal/json" ) // ClientRegistrationMetadata represents the client metadata fields for the DCR POST request (RFC 7591). +// +// Note: URL fields in this struct are validated by validateClientRegistrationURLs +// to prevent XSS attacks. If you add a new URL field, you must also add it to +// that function. type ClientRegistrationMetadata struct { // RedirectURIs is a REQUIRED JSON array of redirection URI strings for use in // redirect-based flows (such as the authorization code grant). @@ -85,6 +89,11 @@ type ClientRegistrationMetadata struct { // SoftwareStatement is an OPTIONAL JWT that asserts client metadata values. // Values in the software statement take precedence over other metadata values. SoftwareStatement string `json:"software_statement,omitempty"` + + // ApplicationType is an OPTIONAL string that indicates the type of application. + // Valid values are "native" and "web". + // If omitted, OIDC-compliant authorization servers default to "web". + ApplicationType string `json:"application_type,omitempty"` } // ClientRegistrationResponse represents the fields returned by the Authorization Server @@ -140,7 +149,7 @@ func (r *ClientRegistrationResponse) UnmarshalJSON(data []byte) error { }{ alias: (*alias)(r), } - if err := json.Unmarshal(data, &aux); err != nil { + if err := internaljson.Unmarshal(data, &aux); err != nil { return err } if aux.ClientIDIssuedAt != 0 { @@ -166,6 +175,10 @@ func (e *ClientRegistrationError) Error() string { return fmt.Sprintf("registration failed: %s (%s)", e.ErrorCode, e.ErrorDescription) } +// maxRegistrationResponseBytes bounds the dynamic client registration response +// body read from the authorization server. +const maxRegistrationResponseBytes = 1 << 20 // 1 MiB + // RegisterClient performs Dynamic Client Registration according to RFC 7591. func RegisterClient(ctx context.Context, registrationEndpoint string, clientMeta *ClientRegistrationMetadata, c *http.Client) (*ClientRegistrationResponse, error) { if registrationEndpoint == "" { @@ -195,25 +208,29 @@ func RegisterClient(ctx context.Context, registrationEndpoint string, clientMeta } defer resp.Body.Close() - body, err := io.ReadAll(resp.Body) + body, err := io.ReadAll(io.LimitReader(resp.Body, maxRegistrationResponseBytes)) if err != nil { return nil, fmt.Errorf("failed to read registration response body: %w", err) } - if resp.StatusCode == http.StatusCreated { + if resp.StatusCode == http.StatusCreated || resp.StatusCode == http.StatusOK { var regResponse ClientRegistrationResponse - if err := json.Unmarshal(body, ®Response); err != nil { + if err := internaljson.Unmarshal(body, ®Response); err != nil { return nil, fmt.Errorf("failed to decode successful registration response: %w (%s)", err, string(body)) } if regResponse.ClientID == "" { return nil, fmt.Errorf("registration response is missing required 'client_id' field") } + // Validate URL fields to prevent XSS attacks (see #526). + if err := validateClientRegistrationURLs(®Response.ClientRegistrationMetadata); err != nil { + return nil, err + } return ®Response, nil } if resp.StatusCode == http.StatusBadRequest { var regError ClientRegistrationError - if err := json.Unmarshal(body, ®Error); err != nil { + if err := internaljson.Unmarshal(body, ®Error); err != nil { return nil, fmt.Errorf("failed to decode registration error response: %w (%s)", err, string(body)) } return nil, ®Error @@ -221,3 +238,33 @@ func RegisterClient(ctx context.Context, registrationEndpoint string, clientMeta return nil, fmt.Errorf("registration failed with status %s: %s", resp.Status, string(body)) } + +// validateClientRegistrationURLs validates all URL fields in ClientRegistrationMetadata +// to ensure they don't use dangerous schemes that could enable XSS attacks. +func validateClientRegistrationURLs(meta *ClientRegistrationMetadata) error { + // Validate redirect URIs + for i, uri := range meta.RedirectURIs { + if err := checkURLScheme(uri); err != nil { + return fmt.Errorf("redirect_uris[%d]: %w", i, err) + } + } + + // Validate other URL fields + urls := []struct { + name string + value string + }{ + {"client_uri", meta.ClientURI}, + {"logo_uri", meta.LogoURI}, + {"tos_uri", meta.TOSURI}, + {"policy_uri", meta.PolicyURI}, + {"jwks_uri", meta.JWKSURI}, + } + + for _, u := range urls { + if err := checkURLScheme(u.value); err != nil { + return fmt.Errorf("%s: %w", u.name, err) + } + } + return nil +} diff --git a/oauthex/dcr_test.go b/oauthex/dcr_test.go index 7c5551e..afff34c 100644 --- a/oauthex/dcr_test.go +++ b/oauthex/dcr_test.go @@ -2,8 +2,6 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. -//go:build mcp_go_client_oauth - package oauthex import ( @@ -72,6 +70,30 @@ func TestRegisterClient(t *testing.T) { clientMeta: &ClientRegistrationMetadata{ClientName: "Test App", RedirectURIs: []string{"http://localhost/cb"}}, wantClientID: "test-client-id", }, + { + name: "Success with 200", + handler: func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + t.Errorf("Expected POST, got %s", r.Method) + } + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatal(err) + } + var receivedMeta ClientRegistrationMetadata + if err := json.Unmarshal(body, &receivedMeta); err != nil { + t.Fatalf("Failed to unmarshal request body: %v", err) + } + if receivedMeta.ClientName != "Test App" { + t.Errorf("Expected ClientName 'Test App', got '%s'", receivedMeta.ClientName) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"client_id":"test-client-id","client_secret":"test-client-secret","client_name":"Test App"}`)) + }, + clientMeta: &ClientRegistrationMetadata{ClientName: "Test App", RedirectURIs: []string{"http://localhost/cb"}}, + wantClientID: "test-client-id", + }, { name: "Missing ClientID in Response", handler: func(w http.ResponseWriter, r *http.Request) { diff --git a/oauthex/discovery_test.go b/oauthex/discovery_test.go new file mode 100644 index 0000000..8b64d0c --- /dev/null +++ b/oauthex/discovery_test.go @@ -0,0 +1,158 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package oauthex + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestCheckHTTPSOrLoopback(t *testing.T) { + tests := []struct { + name string + addr string + wantErr bool + }{ + {name: "empty", addr: "", wantErr: false}, + {name: "https dns allowed", addr: "https://example.com/prm", wantErr: false}, + {name: "cleartext dns disallowed", addr: "http://example.com/prm", wantErr: true}, + {name: "https loopback ip allowed", addr: "http://127.0.0.1:8080/prm", wantErr: false}, + {name: "http localhost allowed", addr: "http://localhost/prm", wantErr: false}, + {name: "https private ip disallowed", addr: "https://10.0.0.1/prm", wantErr: true}, + {name: "link local metadata disallowed", addr: "https://169.254.169.254/", wantErr: true}, + {name: "mapped private address disallowed", addr: "https://[::ffff:10.0.0.1]/prm", wantErr: true}, + {name: "https public ip allowed", addr: "https://8.8.8.8/prm", wantErr: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := checkHTTPSOrLoopback(tt.addr) + if (err != nil) != tt.wantErr { + t.Errorf("checkHTTPSOrLoopback(%q) error = %v, wantErr %v", tt.addr, err, tt.wantErr) + } + }) + } +} + +func TestDiscoveryClientRedirect(t *testing.T) { + newRequestTo := func(u string) *http.Request { + t.Helper() + r, err := http.NewRequest(http.MethodGet, u, nil) + if err != nil { + t.Fatalf("NewRequest(%q): %v", u, err) + } + return r + } + testCases := []struct { + name string + custom *http.Client + req *http.Request + via []*http.Request + wantBlocked bool + }{ + { + name: "redirect to public allowed", + req: newRequestTo("https://8.8.8.8/"), + via: []*http.Request{newRequestTo("https://93.184.216.34/")}, + wantBlocked: false, + }, + { + name: "downgrade blocked", + req: newRequestTo("http://93.184.216.34/"), + via: []*http.Request{newRequestTo("https://93.184.216.34/")}, + wantBlocked: true, + }, + { + name: "redirect to loopback blocked", + req: newRequestTo("https://127.0.0.1/"), + via: []*http.Request{newRequestTo("https://93.184.216.34/")}, + wantBlocked: true, + }, + { + name: "redirect to hostname resolving to loopback blocked", + req: newRequestTo("https://localhost/"), + via: []*http.Request{newRequestTo("https://93.184.216.34/")}, + wantBlocked: true, + }, + { + name: "custom client with no custom check", + custom: &http.Client{}, + req: newRequestTo("https://127.0.0.1/"), + via: []*http.Request{newRequestTo("https://93.184.216.34/")}, + wantBlocked: true, + }, + { + name: "custom client check disable", + custom: &http.Client{ + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return nil + }, + }, + req: newRequestTo("https://127.0.0.1/"), + via: []*http.Request{newRequestTo("https://93.184.216.34/")}, + wantBlocked: false, + }, + { + name: "too many redirects", + req: newRequestTo("https://8.8.8.8/"), + via: func() []*http.Request { + via := make([]*http.Request, maxDiscoveryRedirects) + for i := range via { + via[i] = newRequestTo("https://8.8.8.8/") + } + return via + }(), + wantBlocked: true, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + client := newDiscoveryClient(tc.custom) + gotErr := client.CheckRedirect(tc.req, tc.via) + if gotErr != nil && !tc.wantBlocked { + t.Fatalf("client.CheckRedirect() error = %v", gotErr) + } + if gotErr == nil && tc.wantBlocked { + t.Fatal("client.CheckRedirect() error = nil, want blocked") + } + }) + } +} + +func TestDiscoveryBlocksInternalTarget(t *testing.T) { + ctx := t.Context() + for _, target := range []string{ + "https://169.254.169.254/latest/meta-data/", + "https://10.0.0.1/.well-known/oauth-protected-resource", + } { + t.Run(target, func(t *testing.T) { + _, err := GetProtectedResourceMetadata(ctx, target, target, nil) + if err == nil { + t.Fatalf("GetProtectedResourceMetadata(%q) succeeded, want SSRF rejection", target) + } + if !strings.Contains(err.Error(), "non-public") { + t.Errorf("error = %v, want it to mention a non-public address", err) + } + }) + } +} + +func TestDiscoveryAllowsLoopbackInitial(t *testing.T) { + ctx := t.Context() + h := &fakeResourceHandler{} + server := httptest.NewTLSServer(h) + defer server.Close() + h.installHandlers(server.URL) + + metadataURL := server.URL + "/.well-known/oauth-protected-resource" + prm, err := GetProtectedResourceMetadata(ctx, metadataURL, server.URL, server.Client()) + if err != nil { + t.Fatalf("GetProtectedResourceMetadata on loopback failed: %v", err) + } + if prm == nil { + t.Fatal("nil prm") + } +} diff --git a/oauthex/oauth2.go b/oauthex/oauth2.go index f876be0..0e77149 100644 --- a/oauthex/oauth2.go +++ b/oauthex/oauth2.go @@ -4,35 +4,44 @@ // Package oauthex implements extensions to OAuth2. -//go:build mcp_go_client_oauth - package oauthex import ( "context" - "encoding/json" "fmt" "io" + "mime" + "net" "net/http" + "net/netip" "net/url" "strings" + "sync" + "syscall" + "time" + + "github.com/modelcontextprotocol/go-sdk/internal/json" + "github.com/modelcontextprotocol/go-sdk/internal/util" ) -// prependToPath prepends pre to the path of urlStr. -// When pre is the well-known path, this is the algorithm specified in both RFC 9728 -// section 3.1 and RFC 8414 section 3.1. -func prependToPath(urlStr, pre string) (string, error) { - u, err := url.Parse(urlStr) - if err != nil { - return "", err - } - p := "/" + strings.Trim(pre, "/") - if u.Path != "" { - p += "/" - } +const maxDiscoveryRedirects = 10 - u.Path = p + strings.TrimLeft(u.Path, "/") - return u.String(), nil +// Build this on first use, not when the package loads. Building it checks the +// proxy via http.ProxyFromEnvironment, and net/http only reads the proxy +// environment once and then remembers it for the whole run. If that happened +// at import time, it would lock in the proxy setting before the app had a +// chance to set HTTPS_PROXY. sync.OnceValue waits until first use and still +// builds the transport just once. See #1276. +var defaultDiscoveryTransport = sync.OnceValue(func() http.RoundTripper { + return newDiscoveryTransport(http.DefaultTransport) +}) + +type httpStatusError struct { + StatusCode int +} + +func (e *httpStatusError) Error() string { + return fmt.Sprintf("bad status %d", e.StatusCode) } // getJSON retrieves JSON and unmarshals JSON from the URL, as specified in both @@ -43,21 +52,18 @@ func getJSON[T any](ctx context.Context, c *http.Client, url string, limit int64 if err != nil { return nil, err } - if c == nil { - c = http.DefaultClient - } - res, err := c.Do(req) + res, err := newDiscoveryClient(c).Do(req) if err != nil { return nil, err } defer res.Body.Close() - // Specs require a 200. if res.StatusCode != http.StatusOK { - return nil, fmt.Errorf("bad status %s", res.Status) + return nil, &httpStatusError{StatusCode: res.StatusCode} } - // Specs require application/json. - if ct := res.Header.Get("Content-Type"); ct != "application/json" { + ct := res.Header.Get("Content-Type") + mediaType, _, err := mime.ParseMediaType(ct) + if err != nil || mediaType != "application/json" { return nil, fmt.Errorf("bad content type %q", ct) } @@ -72,6 +78,7 @@ func getJSON[T any](ctx context.Context, c *http.Client, url string, limit int64 // checkURLScheme ensures that its argument is a valid URL with a scheme // that prevents XSS attacks. // See #526. +// Note: a copy of this function exists in auth/extauth/oidc_login.go; keep these in sync. func checkURLScheme(u string) error { if u == "" { return nil @@ -86,3 +93,124 @@ func checkURLScheme(u string) error { } return nil } + +func checkHTTPSOrLoopback(addr string) error { + if addr == "" { + return nil + } + u, err := url.Parse(addr) + if err != nil { + return err + } + if !util.IsLoopback(u.Host) && u.Scheme != "https" { + return fmt.Errorf("URL %q does not use HTTPS or is not a loopback address", addr) + } + if ip, err := netip.ParseAddr(u.Hostname()); err == nil && util.IsPrivateOrReserved(ip) { + return fmt.Errorf("URL %q refers to a non-public address", addr) + } + return nil +} + +func newDiscoveryClient(c *http.Client) *http.Client { + transport := defaultDiscoveryTransport() + if c != nil && c.Transport != nil { + // a caller that has taken over dialing or TLS opts out of the checks + if t, ok := c.Transport.(*http.Transport); ok && t.DialContext == nil && t.DialTLSContext == nil { + transport = newDiscoveryTransport(t) + } else { + transport = c.Transport + } + } + discoveryClient := &http.Client{ + Transport: transport, + + CheckRedirect: func(req *http.Request, via []*http.Request) error { + if c != nil && c.CheckRedirect != nil { + return c.CheckRedirect(req, via) + } + if len(via) >= maxDiscoveryRedirects { + return fmt.Errorf("max redirects exceeded (%d)", maxDiscoveryRedirects) + } + if prev := via[len(via)-1]; prev.URL.Scheme == "https" && req.URL.Scheme != "https" { + return fmt.Errorf("redirect to non-HTTPS scheme %q is a security downgrade", req.URL.Scheme) + } + return checkLoopbackRedirect(req.Context(), req.URL.Hostname()) + }, + } + if c != nil { + discoveryClient.Timeout = c.Timeout + discoveryClient.Jar = c.Jar + } + return discoveryClient +} + +// newDiscoveryTransport creates an http.RoundTripper with extra protections enabled, +// unless the provided RoundTripper is not *http.Transport. +func newDiscoveryTransport(rt http.RoundTripper) http.RoundTripper { + base, ok := rt.(*http.Transport) + if !ok { + return rt + } + // If a proxy is configured Control checks should be disabled, because they will + // basically be running for proxy dials. + if proxyConfigured(base) { + return rt + } + result := base.Clone() + result.DialContext = (&net.Dialer{ + Timeout: 30 * time.Second, + KeepAlive: 30 * time.Second, + Control: func(_, address string, _ syscall.RawConn) error { + host, _, err := net.SplitHostPort(address) + if err != nil { + host = address + } + ip, err := netip.ParseAddr(host) + if err != nil { + return fmt.Errorf("non-ip address: %w", err) + } + if util.IsPrivateOrReserved(ip) { + return fmt.Errorf("non-public ip address: %q", ip) + } + return nil + }, + }).DialContext + return result +} + +func proxyConfigured(t *http.Transport) bool { + if t.Proxy == nil { + return false + } + req, err := http.NewRequest(http.MethodGet, "https://check-proxy-existence.com", nil) + if err != nil { + return false + } + u, err := t.Proxy(req) + return err == nil && u != nil +} + +func checkLoopbackRedirect(ctx context.Context, host string) error { + doCheck := func(ip netip.Addr) error { + if ip.Unmap().IsLoopback() { + return fmt.Errorf("redirect into loopback address %q", host) + } + if util.IsPrivateOrReserved(ip) { + return fmt.Errorf("redirect into non-public address %q", host) + } + return nil + } + if ip, err := netip.ParseAddr(host); err == nil { + return doCheck(ip) + } + ips, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host) + if err != nil { + return fmt.Errorf("cannot resolve redirect host %q: %w", host, err) + } + for _, ip := range ips { + if err := doCheck(ip); err != nil { + return err + } + } + return nil +} diff --git a/oauthex/oauthex.go b/oauthex/oauthex.go index 34ed55b..151da7e 100644 --- a/oauthex/oauthex.go +++ b/oauthex/oauthex.go @@ -4,89 +4,3 @@ // Package oauthex implements extensions to OAuth2. package oauthex - -// ProtectedResourceMetadata is the metadata for an OAuth 2.0 protected resource, -// as defined in section 2 of https://www.rfc-editor.org/rfc/rfc9728.html. -// -// The following features are not supported: -// - additional keys (§2, last sentence) -// - human-readable metadata (§2.1) -// - signed metadata (§2.2) -type ProtectedResourceMetadata struct { - // GENERATED BY GEMINI 2.5. - - // Resource (resource) is the protected resource's resource identifier. - // Required. - Resource string `json:"resource"` - - // AuthorizationServers (authorization_servers) is an optional slice containing a list of - // OAuth authorization server issuer identifiers (as defined in RFC 8414) that can be - // used with this protected resource. - AuthorizationServers []string `json:"authorization_servers,omitempty"` - - // JWKSURI (jwks_uri) is an optional URL of the protected resource's JSON Web Key (JWK) Set - // document. This contains public keys belonging to the protected resource, such as - // signing key(s) that the resource server uses to sign resource responses. - JWKSURI string `json:"jwks_uri,omitempty"` - - // ScopesSupported (scopes_supported) is a recommended slice containing a list of scope - // values (as defined in RFC 6749) used in authorization requests to request access - // to this protected resource. - ScopesSupported []string `json:"scopes_supported,omitempty"` - - // BearerMethodsSupported (bearer_methods_supported) is an optional slice containing - // a list of the supported methods of sending an OAuth 2.0 bearer token to the - // protected resource. Defined values are "header", "body", and "query". - BearerMethodsSupported []string `json:"bearer_methods_supported,omitempty"` - - // ResourceSigningAlgValuesSupported (resource_signing_alg_values_supported) is an optional - // slice of JWS signing algorithms (alg values) supported by the protected - // resource for signing resource responses. - ResourceSigningAlgValuesSupported []string `json:"resource_signing_alg_values_supported,omitempty"` - - // ResourceName (resource_name) is a human-readable name of the protected resource - // intended for display to the end user. It is RECOMMENDED that this field be included. - // This value may be internationalized. - ResourceName string `json:"resource_name,omitempty"` - - // ResourceDocumentation (resource_documentation) is an optional URL of a page containing - // human-readable information for developers using the protected resource. - // This value may be internationalized. - ResourceDocumentation string `json:"resource_documentation,omitempty"` - - // ResourcePolicyURI (resource_policy_uri) is an optional URL of a page containing - // human-readable policy information on how a client can use the data provided. - // This value may be internationalized. - ResourcePolicyURI string `json:"resource_policy_uri,omitempty"` - - // ResourceTOSURI (resource_tos_uri) is an optional URL of a page containing the protected - // resource's human-readable terms of service. This value may be internationalized. - ResourceTOSURI string `json:"resource_tos_uri,omitempty"` - - // TLSClientCertificateBoundAccessTokens (tls_client_certificate_bound_access_tokens) is an - // optional boolean indicating support for mutual-TLS client certificate-bound - // access tokens (RFC 8705). Defaults to false if omitted. - TLSClientCertificateBoundAccessTokens bool `json:"tls_client_certificate_bound_access_tokens,omitempty"` - - // AuthorizationDetailsTypesSupported (authorization_details_types_supported) is an optional - // slice of 'type' values supported by the resource server for the - // 'authorization_details' parameter (RFC 9396). - AuthorizationDetailsTypesSupported []string `json:"authorization_details_types_supported,omitempty"` - - // DPOPSigningAlgValuesSupported (dpop_signing_alg_values_supported) is an optional - // slice of JWS signing algorithms supported by the resource server for validating - // DPoP proof JWTs (RFC 9449). - DPOPSigningAlgValuesSupported []string `json:"dpop_signing_alg_values_supported,omitempty"` - - // DPOPBoundAccessTokensRequired (dpop_bound_access_tokens_required) is an optional boolean - // specifying whether the protected resource always requires the use of DPoP-bound - // access tokens (RFC 9449). Defaults to false if omitted. - DPOPBoundAccessTokensRequired bool `json:"dpop_bound_access_tokens_required,omitempty"` - - // SignedMetadata (signed_metadata) is an optional JWT containing metadata parameters - // about the protected resource as claims. If present, these values take precedence - // over values conveyed in plain JSON. - // TODO:implement. - // Note that §2.2 says it's okay to ignore this. - // SignedMetadata string `json:"signed_metadata,omitempty"` -} diff --git a/oauthex/oauthex_test.go b/oauthex/oauthex_test.go index ca19ace..d028037 100644 --- a/oauthex/oauthex_test.go +++ b/oauthex/oauthex_test.go @@ -2,8 +2,6 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. -//go:build mcp_go_client_oauth - package oauthex import ( diff --git a/oauthex/resource_meta.go b/oauthex/resource_meta.go index bb61f79..7506634 100644 --- a/oauthex/resource_meta.go +++ b/oauthex/resource_meta.go @@ -5,8 +5,6 @@ // This file implements Protected Resource Metadata. // See https://www.rfc-editor.org/rfc/rfc9728.html. -//go:build mcp_go_client_oauth - package oauthex import ( @@ -14,94 +12,99 @@ import ( "errors" "fmt" "net/http" - "net/url" - "path" "strings" "unicode" "github.com/modelcontextprotocol/go-sdk/internal/util" ) -const defaultProtectedResourceMetadataURI = "/.well-known/oauth-protected-resource" +// ProtectedResourceMetadata is the metadata for an OAuth 2.0 protected resource, +// as defined in section 2 of https://www.rfc-editor.org/rfc/rfc9728.html. +// +// The following features are not supported: +// - additional keys (§2, last sentence) +// - human-readable metadata (§2.1) +// - signed metadata (§2.2) +type ProtectedResourceMetadata struct { + // Resource (resource) is the protected resource's resource identifier. + // Required. + Resource string `json:"resource"` -// GetProtectedResourceMetadataFromID issues a GET request to retrieve protected resource -// metadata from a resource server by its ID. -// The resource ID is an HTTPS URL, typically with a host:port and possibly a path. -// For example: -// -// https://example.com/server -// -// This function, following the spec (§3), inserts the default well-known path into the -// URL. In our example, the result would be -// -// https://example.com/.well-known/oauth-protected-resource/server -// -// It then retrieves the metadata at that location using the given client (or the -// default client if nil) and validates its resource field against resourceID. -func GetProtectedResourceMetadataFromID(ctx context.Context, resourceID string, c *http.Client) (_ *ProtectedResourceMetadata, err error) { - defer util.Wrapf(&err, "GetProtectedResourceMetadataFromID(%q)", resourceID) + // AuthorizationServers (authorization_servers) is an optional slice containing a list of + // OAuth authorization server issuer identifiers (as defined in RFC 8414) that can be + // used with this protected resource. + AuthorizationServers []string `json:"authorization_servers,omitempty"` - u, err := url.Parse(resourceID) - if err != nil { - return nil, err - } - // Insert well-known URI into URL. - u.Path = path.Join(defaultProtectedResourceMetadataURI, u.Path) - return getPRM(ctx, u.String(), c, resourceID) + // JWKSURI (jwks_uri) is an optional URL of the protected resource's JSON Web Key (JWK) Set + // document. This contains public keys belonging to the protected resource, such as + // signing key(s) that the resource server uses to sign resource responses. + JWKSURI string `json:"jwks_uri,omitempty"` + + // ScopesSupported (scopes_supported) is a recommended slice containing a list of scope + // values (as defined in RFC 6749) used in authorization requests to request access + // to this protected resource. + ScopesSupported []string `json:"scopes_supported,omitempty"` + + // BearerMethodsSupported (bearer_methods_supported) is an optional slice containing + // a list of the supported methods of sending an OAuth 2.0 bearer token to the + // protected resource. Defined values are "header", "body", and "query". + BearerMethodsSupported []string `json:"bearer_methods_supported,omitempty"` + + // ResourceSigningAlgValuesSupported (resource_signing_alg_values_supported) is an optional + // slice of JWS signing algorithms (alg values) supported by the protected + // resource for signing resource responses. + ResourceSigningAlgValuesSupported []string `json:"resource_signing_alg_values_supported,omitempty"` + + // ResourceName (resource_name) is a human-readable name of the protected resource + // intended for display to the end user. It is RECOMMENDED that this field be included. + // This value may be internationalized. + ResourceName string `json:"resource_name,omitempty"` + + // ResourceDocumentation (resource_documentation) is an optional URL of a page containing + // human-readable information for developers using the protected resource. + // This value may be internationalized. + ResourceDocumentation string `json:"resource_documentation,omitempty"` + + // ResourcePolicyURI (resource_policy_uri) is an optional URL of a page containing + // human-readable policy information on how a client can use the data provided. + // This value may be internationalized. + ResourcePolicyURI string `json:"resource_policy_uri,omitempty"` + + // ResourceTOSURI (resource_tos_uri) is an optional URL of a page containing the protected + // resource's human-readable terms of service. This value may be internationalized. + ResourceTOSURI string `json:"resource_tos_uri,omitempty"` + + // TLSClientCertificateBoundAccessTokens (tls_client_certificate_bound_access_tokens) is an + // optional boolean indicating support for mutual-TLS client certificate-bound + // access tokens (RFC 8705). Defaults to false if omitted. + TLSClientCertificateBoundAccessTokens bool `json:"tls_client_certificate_bound_access_tokens,omitempty"` + + // AuthorizationDetailsTypesSupported (authorization_details_types_supported) is an optional + // slice of 'type' values supported by the resource server for the + // 'authorization_details' parameter (RFC 9396). + AuthorizationDetailsTypesSupported []string `json:"authorization_details_types_supported,omitempty"` + + // DPOPSigningAlgValuesSupported (dpop_signing_alg_values_supported) is an optional + // slice of JWS signing algorithms supported by the resource server for validating + // DPoP proof JWTs (RFC 9449). + DPOPSigningAlgValuesSupported []string `json:"dpop_signing_alg_values_supported,omitempty"` + + // DPOPBoundAccessTokensRequired (dpop_bound_access_tokens_required) is an optional boolean + // specifying whether the protected resource always requires the use of DPoP-bound + // access tokens (RFC 9449). Defaults to false if omitted. + DPOPBoundAccessTokensRequired bool `json:"dpop_bound_access_tokens_required,omitempty"` + + // SignedMetadata (signed_metadata) is an optional JWT containing metadata parameters + // about the protected resource as claims. If present, these values take precedence + // over values conveyed in plain JSON. + // TODO:implement. + // Note that §2.2 says it's okay to ignore this. + // SignedMetadata string `json:"signed_metadata,omitempty"` } -// GetProtectedResourceMetadataFromHeader retrieves protected resource metadata -// using information in the given header, using the given client (or the default -// client if nil). -// It issues a GET request to a URL discovered by parsing the WWW-Authenticate headers in the given request. -// Per RFC 9728 section 3.3, it validates that the resource field of the resulting metadata -// matches the serverURL (the URL that the client used to make the original request to the resource server). -// If there is no metadata URL in the header, it returns nil, nil. -func GetProtectedResourceMetadataFromHeader(ctx context.Context, serverURL string, header http.Header, c *http.Client) (_ *ProtectedResourceMetadata, err error) { - defer util.Wrapf(&err, "GetProtectedResourceMetadataFromHeader") - headers := header[http.CanonicalHeaderKey("WWW-Authenticate")] - if len(headers) == 0 { - return nil, nil - } - cs, err := ParseWWWAuthenticate(headers) - if err != nil { - return nil, err - } - metadataURL := ResourceMetadataURL(cs) - if metadataURL == "" { - return nil, nil - } - return getPRM(ctx, metadataURL, c, serverURL) -} - -// getPRM makes a GET request to the given URL, and validates the response. -// As part of the validation, it compares the returned resource field to wantResource. -func getPRM(ctx context.Context, purl string, c *http.Client, wantResource string) (*ProtectedResourceMetadata, error) { - if !strings.HasPrefix(strings.ToUpper(purl), "HTTPS://") { - return nil, fmt.Errorf("resource URL %q does not use HTTPS", purl) - } - prm, err := getJSON[ProtectedResourceMetadata](ctx, c, purl, 1<<20) - if err != nil { - return nil, err - } - // Validate the Resource field (see RFC 9728, section 3.3). - if prm.Resource != wantResource { - return nil, fmt.Errorf("got metadata resource %q, want %q", prm.Resource, wantResource) - } - // Validate the authorization server URLs to prevent XSS attacks (see #526). - for _, u := range prm.AuthorizationServers { - if err := checkURLScheme(u); err != nil { - return nil, err - } - } - return prm, nil -} - -// challenge represents a single authentication challenge from a WWW-Authenticate header. +// Challenge represents a single authentication challenge from a WWW-Authenticate header. // As per RFC 9110, Section 11.6.1, a challenge consists of a scheme and optional parameters. -type challenge struct { - // GENERATED BY GEMINI 2.5. - // +type Challenge struct { // Scheme is the authentication scheme (e.g., "Bearer", "Basic"). // It is case-insensitive. A parsed value will always be lower-case. Scheme string @@ -110,24 +113,46 @@ type challenge struct { Params map[string]string } -// ResourceMetadataURL returns a resource metadata URL from the given challenges, -// or the empty string if there is none. -func ResourceMetadataURL(cs []challenge) string { - for _, c := range cs { - if u := c.Params["resource_metadata"]; u != "" { - return u +// GetProtectedResourceMetadata issues a GET request to retrieve protected resource +// metadata from a resource server. +// The metadataURL is typically a URL with a host:port and possibly a path. +// The resourceURL is the resource URI the metadataURL is for. +// The following checks are performed: +// - The metadataURL must use HTTPS or be a local address. +// - The resource field of the resulting metadata must match the resourceURL. +// - The authorization_servers field of the resulting metadata is checked for dangerous URL schemes. +func GetProtectedResourceMetadata(ctx context.Context, metadataURL, resourceURL string, c *http.Client) (_ *ProtectedResourceMetadata, err error) { + defer util.Wrapf(&err, "GetProtectedResourceMetadata(%q)", metadataURL) + // Only allow HTTP for local addresses (testing or development purposes). + if err := checkHTTPSOrLoopback(metadataURL); err != nil { + return nil, fmt.Errorf("metadataURL: %v", err) + } + prm, err := getJSON[ProtectedResourceMetadata](ctx, c, metadataURL, 1<<20) + if err != nil { + return nil, err + } + // Validate the Resource field (see RFC 9728, section 3.3). + if prm.Resource != resourceURL { + return nil, fmt.Errorf("got metadata resource %q, want %q", prm.Resource, resourceURL) + } + // Validate the authorization server URLs to prevent XSS attacks (see #526). + for i, u := range prm.AuthorizationServers { + if err := checkURLScheme(u); err != nil { + return nil, fmt.Errorf("authorization_servers[%d]: %v", i, err) + } + if err := checkHTTPSOrLoopback(u); err != nil { + return nil, fmt.Errorf("authorization_servers[%d]: %v", i, err) } } - return "" + return prm, nil } // ParseWWWAuthenticate parses a WWW-Authenticate header string. // The header format is defined in RFC 9110, Section 11.6.1, and can contain // one or more challenges, separated by commas. // It returns a slice of challenges or an error if one of the headers is malformed. -func ParseWWWAuthenticate(headers []string) ([]challenge, error) { - // GENERATED BY GEMINI 2.5 (human-tweaked) - var challenges []challenge +func ParseWWWAuthenticate(headers []string) ([]Challenge, error) { + var challenges []Challenge for _, h := range headers { challengeStrings, err := splitChallenges(h) if err != nil { @@ -151,7 +176,6 @@ func ParseWWWAuthenticate(headers []string) ([]challenge, error) { // It correctly handles commas within quoted strings and distinguishes between // commas separating auth-params and commas separating challenges. func splitChallenges(header string) ([]string, error) { - // GENERATED BY GEMINI 2.5. var challenges []string inQuotes := false start := 0 @@ -195,15 +219,14 @@ func splitChallenges(header string) ([]string, error) { // parseSingleChallenge parses a string containing exactly one challenge. // challenge = auth-scheme [ 1*SP ( token68 / #auth-param ) ] -func parseSingleChallenge(s string) (challenge, error) { - // GENERATED BY GEMINI 2.5, human-tweaked. +func parseSingleChallenge(s string) (Challenge, error) { s = strings.TrimSpace(s) if s == "" { - return challenge{}, errors.New("empty challenge string") + return Challenge{}, errors.New("empty challenge string") } scheme, paramsStr, found := strings.Cut(s, " ") - c := challenge{Scheme: strings.ToLower(scheme)} + c := Challenge{Scheme: strings.ToLower(scheme)} if !found { return c, nil } @@ -215,7 +238,7 @@ func parseSingleChallenge(s string) (challenge, error) { // Find the end of the parameter key. keyEnd := strings.Index(paramsStr, "=") if keyEnd <= 0 { - return challenge{}, fmt.Errorf("malformed auth parameter: expected key=value, but got %q", paramsStr) + return Challenge{}, fmt.Errorf("malformed auth parameter: expected key=value, but got %q", paramsStr) } key := strings.TrimSpace(paramsStr[:keyEnd]) @@ -243,7 +266,7 @@ func parseSingleChallenge(s string) (challenge, error) { // A quoted string must be terminated. if i == len(paramsStr) { - return challenge{}, fmt.Errorf("unterminated quoted string in auth parameter") + return Challenge{}, fmt.Errorf("unterminated quoted string in auth parameter") } value = valBuilder.String() @@ -261,7 +284,7 @@ func parseSingleChallenge(s string) (challenge, error) { } } if value == "" { - return challenge{}, fmt.Errorf("no value for auth param %q", key) + return Challenge{}, fmt.Errorf("no value for auth param %q", key) } // Per RFC 9110, parameter keys are case-insensitive. @@ -272,10 +295,10 @@ func parseSingleChallenge(s string) (challenge, error) { paramsStr = strings.TrimSpace(paramsStr[1:]) } else if paramsStr != "" { // If there's content but it's not a new parameter, the format is wrong. - return challenge{}, fmt.Errorf("malformed auth parameter: expected comma after value, but got %q", paramsStr) + return Challenge{}, fmt.Errorf("malformed auth parameter: expected comma after value, but got %q", paramsStr) } } // Per RFC 9110, the scheme is case-insensitive. - return challenge{Scheme: strings.ToLower(scheme), Params: params}, nil + return Challenge{Scheme: strings.ToLower(scheme), Params: params}, nil } diff --git a/oauthex/oauth2_test.go b/oauthex/resource_meta_test.go similarity index 68% rename from oauthex/oauth2_test.go rename to oauthex/resource_meta_test.go index 08d2d31..81e1ebe 100644 --- a/oauthex/oauth2_test.go +++ b/oauthex/resource_meta_test.go @@ -2,14 +2,10 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. -//go:build mcp_go_client_oauth - package oauthex import ( - "context" "encoding/json" - "fmt" "net/http" "net/http/httptest" "reflect" @@ -82,13 +78,13 @@ func TestParseSingleChallenge(t *testing.T) { tests := []struct { name string input string - want challenge + want Challenge wantErr bool }{ { name: "scheme only", input: "Basic", - want: challenge{ + want: Challenge{ Scheme: "basic", }, wantErr: false, @@ -96,7 +92,7 @@ func TestParseSingleChallenge(t *testing.T) { { name: "scheme with one quoted param", input: `Bearer realm="example.com"`, - want: challenge{ + want: Challenge{ Scheme: "bearer", Params: map[string]string{"realm": "example.com"}, }, @@ -105,7 +101,7 @@ func TestParseSingleChallenge(t *testing.T) { { name: "scheme with one unquoted param", input: `Bearer realm=example.com`, - want: challenge{ + want: Challenge{ Scheme: "bearer", Params: map[string]string{"realm": "example.com"}, }, @@ -114,7 +110,7 @@ func TestParseSingleChallenge(t *testing.T) { { name: "scheme with multiple params", input: `Bearer realm="example", error="invalid_token", error_description="The token expired"`, - want: challenge{ + want: Challenge{ Scheme: "bearer", Params: map[string]string{ "realm": "example", @@ -127,7 +123,7 @@ func TestParseSingleChallenge(t *testing.T) { { name: "scheme with multiple unquoted params", input: `Bearer realm=example, error=invalid_token, error_description=The token expired`, - want: challenge{ + want: Challenge{ Scheme: "bearer", Params: map[string]string{ "realm": "example", @@ -140,7 +136,7 @@ func TestParseSingleChallenge(t *testing.T) { { name: "case-insensitive scheme and keys", input: `BEARER ReAlM="example"`, - want: challenge{ + want: Challenge{ Scheme: "bearer", Params: map[string]string{"realm": "example"}, }, @@ -149,7 +145,7 @@ func TestParseSingleChallenge(t *testing.T) { { name: "param with escaped quote", input: `Bearer realm="example \"foo\" bar"`, - want: challenge{ + want: Challenge{ Scheme: "bearer", Params: map[string]string{"realm": `example "foo" bar`}, }, @@ -158,7 +154,7 @@ func TestParseSingleChallenge(t *testing.T) { { name: "param without quotes (token)", input: "Bearer realm=example.com", - want: challenge{ + want: Challenge{ Scheme: "bearer", Params: map[string]string{"realm": "example.com"}, }, @@ -206,21 +202,14 @@ func TestParseSingleChallenge(t *testing.T) { } func TestGetProtectedResourceMetadata(t *testing.T) { - ctx := context.Background() - t.Run("FromHeader", func(t *testing.T) { - h := &fakeResourceHandler{serveWWWAuthenticate: true} + ctx := t.Context() + t.Run("Success", func(t *testing.T) { + h := &fakeResourceHandler{} server := httptest.NewTLSServer(h) h.installHandlers(server.URL) client := server.Client() - serverURL := server.URL + "/resource" - res, err := client.Get(serverURL) - if err != nil { - t.Fatal(err) - } - if res.StatusCode != http.StatusUnauthorized { - t.Fatal("want unauth") - } - prm, err := GetProtectedResourceMetadataFromHeader(ctx, serverURL, res.Header, client) + metadataURL := server.URL + "/.well-known/oauth-protected-resource" + prm, err := GetProtectedResourceMetadata(ctx, metadataURL, server.URL, client) if err != nil { t.Fatal(err) } @@ -228,53 +217,13 @@ func TestGetProtectedResourceMetadata(t *testing.T) { t.Fatal("nil prm") } }) - t.Run("FromID", func(t *testing.T) { - h := &fakeResourceHandler{serveWWWAuthenticate: false} + t.Run("RejectsIncorrectResource", func(t *testing.T) { + h := &fakeResourceHandler{resourceOverride: "https://attacker.com/evil"} server := httptest.NewTLSServer(h) h.installHandlers(server.URL) client := server.Client() - prm, err := GetProtectedResourceMetadataFromID(ctx, server.URL, client) - if err != nil { - t.Fatal(err) - } - if prm == nil { - t.Fatal("nil prm") - } - }) - // Test that metadata URL and resource identifier are properly distinguished (issue #560). - t.Run("FromHeaderValidatesAgainstServerURL", func(t *testing.T) { - h := &fakeResourceHandler{serveWWWAuthenticate: true} - server := httptest.NewTLSServer(h) - h.installHandlers(server.URL) - client := server.Client() - serverURL := server.URL + "/resource" - res, err := client.Get(serverURL) - if err != nil { - t.Fatal(err) - } - // This should succeed because we validate against serverURL, not metadataURL - prm, err := GetProtectedResourceMetadataFromHeader(ctx, serverURL, res.Header, client) - if err != nil { - t.Fatalf("Expected validation to succeed, got error: %v", err) - } - if prm == nil { - t.Fatal("Expected non-nil prm") - } - if prm.Resource != serverURL { - t.Errorf("Expected resource %q, got %q", serverURL, prm.Resource) - } - }) - t.Run("FromHeaderRejectsImpersonation", func(t *testing.T) { - h := &fakeResourceHandler{serveWWWAuthenticate: true, resourceOverride: "https://attacker.com/evil"} - server := httptest.NewTLSServer(h) - h.installHandlers(server.URL) - client := server.Client() - serverURL := server.URL + "/resource" - res, err := client.Get(serverURL) - if err != nil { - t.Fatal(err) - } - prm, err := GetProtectedResourceMetadataFromHeader(ctx, serverURL, res.Header, client) + metadataURL := server.URL + "/.well-known/oauth-protected-resource" + prm, err := GetProtectedResourceMetadata(ctx, metadataURL, server.URL, client) if err == nil { t.Fatal("Expected validation error for mismatched resource, got nil") } @@ -286,28 +235,17 @@ func TestGetProtectedResourceMetadata(t *testing.T) { type fakeResourceHandler struct { http.ServeMux - serveWWWAuthenticate bool - resourceOverride string // If set, use this instead of correct resource (for testing validation) + resourceOverride string // If set, use this instead of correct resource (for testing validation) } func (h *fakeResourceHandler) installHandlers(serverURL string) { path := "/.well-known/oauth-protected-resource" - url := serverURL + path - h.Handle("GET /resource", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if h.serveWWWAuthenticate { - w.Header().Set("WWW-Authenticate", fmt.Sprintf(`Bearer resource_metadata="%s"`, url)) - } - w.WriteHeader(http.StatusUnauthorized) - })) h.Handle("GET "+path, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") // Per RFC 9728 section 3.3, the resource field should contain the actual resource identifier, // which is the URL the client uses to access the resource (serverURL + "/resource" for WWW-Authenticate case). - // For the FromID test case, it's just the serverURL. + // For the well-known URL test case, it's just the serverURL. resource := serverURL - if h.serveWWWAuthenticate { - resource = serverURL + "/resource" - } // Allow testing with custom resource values (e.g., impersonation attacks). if h.resourceOverride != "" { resource = h.resourceOverride diff --git a/oauthex/token_exchange.go b/oauthex/token_exchange.go new file mode 100644 index 0000000..0c7b24e --- /dev/null +++ b/oauthex/token_exchange.go @@ -0,0 +1,183 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +// This file implements Token Exchange (RFC 8693) for Enterprise Managed Authorization. +// See https://datatracker.ietf.org/doc/html/rfc8693 + +package oauthex + +import ( + "context" + "fmt" + "net/http" + "strings" + + "golang.org/x/oauth2" +) + +// Token type identifiers defined by RFC 8693 and SEP-990. +const ( + // TokenTypeIDToken is the URN for OpenID Connect ID Tokens. + TokenTypeIDToken = "urn:ietf:params:oauth:token-type:id_token" + + // TokenTypeSAML2 is the URN for SAML 2.0 assertions. + TokenTypeSAML2 = "urn:ietf:params:oauth:token-type:saml2" + + // TokenTypeIDJAG is the URN for Identity Assertion JWT Authorization Grants. + // This is the token type returned by IdP during token exchange for SEP-990. + TokenTypeIDJAG = "urn:ietf:params:oauth:token-type:id-jag" + + // GrantTypeTokenExchange is the grant type for RFC 8693 token exchange. + GrantTypeTokenExchange = "urn:ietf:params:oauth:grant-type:token-exchange" +) + +// TokenExchangeRequest represents a Token Exchange request per RFC 8693. +// This is used for Enterprise Managed Authorization (SEP-990) where an MCP Client +// exchanges an ID Token from an enterprise IdP for an ID-JAG that can be used +// to obtain an access token from an MCP Server's authorization server. +type TokenExchangeRequest struct { + // RequestedTokenType indicates the type of security token being requested. + // For SEP-990, this MUST be TokenTypeIDJAG. + RequestedTokenType string + + // Audience is the logical name of the target service where the client + // intends to use the requested token. For SEP-990, this MUST be the + // Issuer URL of the MCP Server's authorization server. + Audience string + + // Resource is the physical location or identifier of the target resource. + // For SEP-990, this MUST be the RFC9728 Resource Identifier of the MCP Server. + Resource string + + // Scope is a list of space-separated scopes for the requested token. + // This is OPTIONAL per RFC 8693 but commonly used in SEP-990. + Scope []string + + // SubjectToken is the security token that represents the identity of the + // party on behalf of whom the request is being made. For SEP-990, this is + // typically an OpenID Connect ID Token. + SubjectToken string + + // SubjectTokenType is the type of the security token in SubjectToken. + // For SEP-990 with OIDC, this MUST be TokenTypeIDToken. + SubjectTokenType string +} + +// ExchangeToken performs a token exchange request per RFC 8693 for Enterprise +// Managed Authorization (SEP-990). It exchanges an identity assertion (typically +// an ID Token) for an Identity Assertion JWT Authorization Grant (ID-JAG) that +// can be used to obtain an access token from an MCP Server. +// +// The tokenEndpoint parameter should be the IdP's token endpoint (typically +// obtained from the IdP's authorization server metadata). +// +// Returns an oauth2.Token where: +// - Extra("issued_token_type") contains the type of the issued token (e.g., TokenTypeIDJAG) +// - AccessToken contains the ID-JAG JWT (despite the name, this is not an OAuth access token) +// - TokenType is typically "N_A" for SEP-990 +// - Extra("scope") may contain the scope if different from the request +// - Expiry is when the token expires +func ExchangeToken( + ctx context.Context, + tokenEndpoint string, + req *TokenExchangeRequest, + clientCreds *ClientCredentials, + httpClient *http.Client, +) (*oauth2.Token, error) { + if tokenEndpoint == "" { + return nil, fmt.Errorf("token endpoint is required") + } + if req == nil { + return nil, fmt.Errorf("token exchange request is required") + } + if clientCreds == nil { + return nil, fmt.Errorf("client credentials are required") + } + if err := clientCreds.Validate(); err != nil { + return nil, fmt.Errorf("invalid client credentials: %w", err) + } + + // Validate required fields per SEP-990 Section 4. + if req.RequestedTokenType == "" { + return nil, fmt.Errorf("requested_token_type is required") + } + if req.Audience == "" { + return nil, fmt.Errorf("audience is required") + } + if req.Resource == "" { + return nil, fmt.Errorf("resource is required") + } + if req.SubjectToken == "" { + return nil, fmt.Errorf("subject_token is required") + } + if req.SubjectTokenType == "" { + return nil, fmt.Errorf("subject_token_type is required") + } + + // Validate URL schemes to prevent XSS attacks (see #526). + if err := checkURLScheme(tokenEndpoint); err != nil { + return nil, fmt.Errorf("invalid token endpoint: %w", err) + } + if err := checkURLScheme(req.Audience); err != nil { + return nil, fmt.Errorf("invalid audience: %w", err) + } + if err := checkURLScheme(req.Resource); err != nil { + return nil, fmt.Errorf("invalid resource: %w", err) + } + + // Per RFC 6749 Section 3.2, parameters sent without a value (like the empty + // "code" parameter) MUST be treated as if they were omitted from the request. + // The oauth2 library's Exchange method sends an empty code, but compliant + // servers should ignore it. + cfg := &oauth2.Config{ + ClientID: clientCreds.ClientID, + Endpoint: oauth2.Endpoint{ + TokenURL: tokenEndpoint, + AuthStyle: oauth2.AuthStyleInParams, + }, + } + // Set ClientSecret if ClientSecretAuth is configured. + if clientCreds.ClientSecretAuth != nil { + cfg.ClientSecret = clientCreds.ClientSecretAuth.ClientSecret + } + + // Use custom HTTP client if provided. + if httpClient == nil { + httpClient = http.DefaultClient + } + ctxWithClient := context.WithValue(ctx, oauth2.HTTPClient, httpClient) + + // Build token exchange parameters per RFC 8693. + opts := []oauth2.AuthCodeOption{ + oauth2.SetAuthURLParam("grant_type", GrantTypeTokenExchange), + oauth2.SetAuthURLParam("requested_token_type", req.RequestedTokenType), + oauth2.SetAuthURLParam("audience", req.Audience), + oauth2.SetAuthURLParam("resource", req.Resource), + oauth2.SetAuthURLParam("subject_token", req.SubjectToken), + oauth2.SetAuthURLParam("subject_token_type", req.SubjectTokenType), + } + if len(req.Scope) > 0 { + opts = append(opts, oauth2.SetAuthURLParam("scope", strings.Join(req.Scope, " "))) + } + + // Exchange with token exchange grant type. + // SetAuthURLParam overrides the default grant_type and adds all required parameters. + token, err := cfg.Exchange( + ctxWithClient, + "", // empty code - per RFC 6749 Section 3.2, empty params should be ignored + opts..., + ) + if err != nil { + return nil, fmt.Errorf("token exchange request failed: %w", err) + } + + // Validate that issued_token_type is present in the response. + // The oauth2 library stores additional response fields in Extra. + issuedTokenType, _ := token.Extra("issued_token_type").(string) + if issuedTokenType == "" { + return nil, fmt.Errorf("response missing required field: issued_token_type") + } + + return token, nil +} diff --git a/oauthex/token_exchange_test.go b/oauthex/token_exchange_test.go new file mode 100644 index 0000000..eb026ff --- /dev/null +++ b/oauthex/token_exchange_test.go @@ -0,0 +1,237 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package oauthex + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +// TestExchangeToken tests the basic token exchange flow. +func TestExchangeToken(t *testing.T) { + // Create a test IdP server that implements token exchange + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Verify request method and content type + if r.Method != http.MethodPost { + t.Errorf("expected POST request, got %s", r.Method) + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + + contentType := r.Header.Get("Content-Type") + if contentType != "application/x-www-form-urlencoded" { + t.Errorf("expected application/x-www-form-urlencoded, got %s", contentType) + http.Error(w, "invalid content type", http.StatusBadRequest) + return + } + + // Parse form data + if err := r.ParseForm(); err != nil { + http.Error(w, "failed to parse form", http.StatusBadRequest) + return + } + + // Verify required parameters per SEP-990 Section 4 + grantType := r.FormValue("grant_type") + if grantType != GrantTypeTokenExchange { + t.Errorf("expected grant_type %s, got %s", GrantTypeTokenExchange, grantType) + writeErrorResponse(w, "invalid_grant", "invalid grant_type") + return + } + + requestedTokenType := r.FormValue("requested_token_type") + if requestedTokenType != TokenTypeIDJAG { + t.Errorf("expected requested_token_type %s, got %s", TokenTypeIDJAG, requestedTokenType) + writeErrorResponse(w, "invalid_request", "invalid requested_token_type") + return + } + + audience := r.FormValue("audience") + if audience == "" { + t.Error("audience is required") + writeErrorResponse(w, "invalid_request", "missing audience") + return + } + + resource := r.FormValue("resource") + if resource == "" { + t.Error("resource is required") + writeErrorResponse(w, "invalid_request", "missing resource") + return + } + + subjectToken := r.FormValue("subject_token") + if subjectToken == "" { + t.Error("subject_token is required") + writeErrorResponse(w, "invalid_request", "missing subject_token") + return + } + + subjectTokenType := r.FormValue("subject_token_type") + if subjectTokenType != TokenTypeIDToken { + t.Errorf("expected subject_token_type %s, got %s", TokenTypeIDToken, subjectTokenType) + writeErrorResponse(w, "invalid_request", "invalid subject_token_type") + return + } + + // Verify client authentication + clientID := r.FormValue("client_id") + clientSecret := r.FormValue("client_secret") + if clientID == "" || clientSecret == "" { + t.Error("client authentication required") + writeErrorResponse(w, "invalid_client", "client authentication failed") + return + } + + if clientID != "test-client-id" || clientSecret != "test-client-secret" { + t.Error("invalid client credentials") + writeErrorResponse(w, "invalid_client", "invalid credentials") + return + } + + // Return successful token exchange response per SEP-990 Section 4.2 + resp := map[string]any{ + "issued_token_type": TokenTypeIDJAG, + "access_token": "fake-id-jag-token", + "token_type": "N_A", + "scope": r.FormValue("scope"), + "expires_in": 300, + } + + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Cache-Control", "no-store") + w.Header().Set("Pragma", "no-cache") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + // Test successful token exchange + t.Run("successful exchange", func(t *testing.T) { + req := &TokenExchangeRequest{ + RequestedTokenType: TokenTypeIDJAG, + Audience: "https://auth.mcpserver.example/", + Resource: "https://mcp.mcpserver.example/", + Scope: []string{"read", "write"}, + SubjectToken: "fake-id-token", + SubjectTokenType: TokenTypeIDToken, + } + + token, err := ExchangeToken( + context.Background(), + server.URL, + req, + &ClientCredentials{ + ClientID: "test-client-id", + ClientSecretAuth: &ClientSecretAuth{ + ClientSecret: "test-client-secret", + }, + }, + server.Client(), + ) + + if err != nil { + t.Fatalf("ExchangeToken failed: %v", err) + } + + issuedTokenType, ok := token.Extra("issued_token_type").(string) + if !ok || issuedTokenType != TokenTypeIDJAG { + t.Errorf("expected issued_token_type %s, got %s", TokenTypeIDJAG, issuedTokenType) + } + + if token.AccessToken != "fake-id-jag-token" { + t.Errorf("expected access_token 'fake-id-jag-token', got %s", token.AccessToken) + } + + if token.TokenType != "N_A" { + t.Errorf("expected token_type 'N_A', got %s", token.TokenType) + } + + scope, ok := token.Extra("scope").(string) + if !ok || scope != "read write" { + t.Errorf("expected scope 'read write', got %s", scope) + } + + // expires_in should be available in Extra + expiresIn, ok := token.Extra("expires_in").(float64) + if !ok || int(expiresIn) != 300 { + t.Errorf("expected expires_in 300, got %v", expiresIn) + } + }) + + // Test missing required fields + t.Run("missing audience", func(t *testing.T) { + req := &TokenExchangeRequest{ + RequestedTokenType: TokenTypeIDJAG, + Resource: "https://mcp.mcpserver.example/", + SubjectToken: "fake-id-token", + SubjectTokenType: TokenTypeIDToken, + } + + _, err := ExchangeToken( + context.Background(), + server.URL, + req, + &ClientCredentials{ + ClientID: "test-client-id", + ClientSecretAuth: &ClientSecretAuth{ + ClientSecret: "test-client-secret", + }, + }, + server.Client(), + ) + + if err == nil { + t.Error("expected error for missing audience, got nil") + } + }) + + // Test invalid URL schemes + t.Run("invalid audience URL scheme", func(t *testing.T) { + req := &TokenExchangeRequest{ + RequestedTokenType: TokenTypeIDJAG, + Audience: "javascript:alert(1)", + Resource: "https://mcp.mcpserver.example/", + SubjectToken: "fake-id-token", + SubjectTokenType: TokenTypeIDToken, + } + + _, err := ExchangeToken( + context.Background(), + server.URL, + req, + &ClientCredentials{ + ClientID: "test-client-id", + ClientSecretAuth: &ClientSecretAuth{ + ClientSecret: "test-client-secret", + }, + }, + server.Client(), + ) + + if err == nil { + t.Error("expected error for invalid audience URL scheme, got nil") + } + }) +} + +// writeErrorResponse writes an OAuth 2.0 error response per RFC 6749 Section 5.2. +func writeErrorResponse(w http.ResponseWriter, errorCode, errorDescription string) { + errResp := struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description,omitempty"` + }{ + Error: errorCode, + ErrorDescription: errorDescription, + } + + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Cache-Control", "no-store") + w.WriteHeader(http.StatusBadRequest) + json.NewEncoder(w).Encode(errResp) +} diff --git a/oauthex/url_scheme_test.go b/oauthex/url_scheme_test.go new file mode 100644 index 0000000..f8c44a3 --- /dev/null +++ b/oauthex/url_scheme_test.go @@ -0,0 +1,315 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package oauthex + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// TestCheckURLScheme tests the checkURLScheme function directly. +func TestCheckURLScheme(t *testing.T) { + tests := []struct { + name string + url string + wantErr bool + }{ + // Valid schemes + {"empty string", "", false}, + {"https url", "https://example.com/path", false}, + {"http url", "http://example.com/path", false}, + {"custom scheme", "myapp://callback", false}, + + // Dangerous schemes that should be blocked + {"javascript scheme", "javascript:alert('xss')", true}, + {"javascript uppercase", "JAVASCRIPT:alert('xss')", true}, + {"javascript mixed case", "JaVaScRiPt:alert('xss')", true}, + {"data scheme", "data:text/html,", true}, + {"data uppercase", "DATA:text/html,test", true}, + {"vbscript scheme", "vbscript:msgbox('xss')", true}, + {"vbscript uppercase", "VBSCRIPT:test", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := checkURLScheme(tt.url) + if (err != nil) != tt.wantErr { + t.Errorf("checkURLScheme(%q): got err %v, want err %v", tt.url, err != nil, tt.wantErr) + } + }) + } +} + +// TestValidateAuthServerMetaURLs tests validation of AuthServerMeta URL fields. +func TestValidateAuthServerMetaURLs(t *testing.T) { + validMeta := &AuthServerMeta{ + Issuer: "https://auth.example.com", + AuthorizationEndpoint: "https://auth.example.com/authorize", + TokenEndpoint: "https://auth.example.com/token", + JWKSURI: "https://auth.example.com/.well-known/jwks.json", + RegistrationEndpoint: "https://auth.example.com/register", + ServiceDocumentation: "https://docs.example.com", + OpPolicyURI: "https://example.com/policy", + OpTOSURI: "https://example.com/tos", + RevocationEndpoint: "https://auth.example.com/revoke", + IntrospectionEndpoint: "https://auth.example.com/introspect", + } + + t.Run("valid metadata", func(t *testing.T) { + if err := validateAuthServerMetaURLs(validMeta); err != nil { + t.Errorf("validateAuthServerMetaURLs(): got err %v, want nil", err) + } + }) + + // Test each URL field with a dangerous scheme + dangerousFields := []struct { + name string + setField func(*AuthServerMeta) + }{ + {"authorization_endpoint", func(m *AuthServerMeta) { m.AuthorizationEndpoint = "javascript:alert(1)" }}, + {"token_endpoint", func(m *AuthServerMeta) { m.TokenEndpoint = "javascript:alert(1)" }}, + {"jwks_uri", func(m *AuthServerMeta) { m.JWKSURI = "data:text/html,test" }}, + {"registration_endpoint", func(m *AuthServerMeta) { m.RegistrationEndpoint = "vbscript:test" }}, + {"service_documentation", func(m *AuthServerMeta) { m.ServiceDocumentation = "javascript:void(0)" }}, + {"op_policy_uri", func(m *AuthServerMeta) { m.OpPolicyURI = "javascript:x" }}, + {"op_tos_uri", func(m *AuthServerMeta) { m.OpTOSURI = "data:,test" }}, + {"revocation_endpoint", func(m *AuthServerMeta) { m.RevocationEndpoint = "javascript:1" }}, + {"introspection_endpoint", func(m *AuthServerMeta) { m.IntrospectionEndpoint = "javascript:2" }}, + } + + for _, tt := range dangerousFields { + t.Run("dangerous "+tt.name, func(t *testing.T) { + // Copy valid metadata + meta := *validMeta + // Set one field to a dangerous value + tt.setField(&meta) + + err := validateAuthServerMetaURLs(&meta) + if err == nil { + t.Errorf("validateAuthServerMetaURLs(): got nil error, want error for dangerous %s", tt.name) + } else if !strings.Contains(err.Error(), tt.name) { + t.Errorf("validateAuthServerMetaURLs(): got error %v, want error containing %q", err, tt.name) + } + }) + } + + t.Run("empty optional fields are valid", func(t *testing.T) { + meta := &AuthServerMeta{ + Issuer: "https://auth.example.com", + AuthorizationEndpoint: "https://auth.example.com/authorize", + TokenEndpoint: "https://auth.example.com/token", + JWKSURI: "https://auth.example.com/.well-known/jwks.json", + // All optional fields left empty + } + if err := validateAuthServerMetaURLs(meta); err != nil { + t.Errorf("validateAuthServerMetaURLs(): got err %v, want nil", err) + } + }) +} + +// TestValidateClientRegistrationURLs tests validation of ClientRegistrationMetadata URL fields. +func TestValidateClientRegistrationURLs(t *testing.T) { + validMeta := &ClientRegistrationMetadata{ + RedirectURIs: []string{"https://app.example.com/callback", "myapp://callback"}, + ClientURI: "https://example.com", + LogoURI: "https://example.com/logo.png", + TOSURI: "https://example.com/tos", + PolicyURI: "https://example.com/policy", + JWKSURI: "https://example.com/.well-known/jwks.json", + } + + t.Run("valid metadata", func(t *testing.T) { + if err := validateClientRegistrationURLs(validMeta); err != nil { + t.Errorf("validateClientRegistrationURLs(): got err %v, want nil", err) + } + }) + + t.Run("dangerous redirect_uri", func(t *testing.T) { + meta := *validMeta + meta.RedirectURIs = []string{"https://safe.com/cb", "javascript:alert(1)"} + + err := validateClientRegistrationURLs(&meta) + if err == nil { + t.Error("validateClientRegistrationURLs(): got nil error, want error for dangerous redirect_uri") + } else if !strings.Contains(err.Error(), "redirect_uris[1]") { + t.Errorf("validateClientRegistrationURLs(): got error %v, want error containing \"redirect_uris[1]\"", err) + } + }) + + // Test each URL field with a dangerous scheme + dangerousFields := []struct { + name string + setField func(*ClientRegistrationMetadata) + }{ + {"client_uri", func(m *ClientRegistrationMetadata) { m.ClientURI = "javascript:alert(1)" }}, + {"logo_uri", func(m *ClientRegistrationMetadata) { m.LogoURI = "data:image/svg," }}, + {"tos_uri", func(m *ClientRegistrationMetadata) { m.TOSURI = "vbscript:test" }}, + {"policy_uri", func(m *ClientRegistrationMetadata) { m.PolicyURI = "javascript:void(0)" }}, + {"jwks_uri", func(m *ClientRegistrationMetadata) { m.JWKSURI = "data:application/json,{}" }}, + } + + for _, tt := range dangerousFields { + t.Run("dangerous "+tt.name, func(t *testing.T) { + meta := *validMeta + tt.setField(&meta) + + err := validateClientRegistrationURLs(&meta) + if err == nil { + t.Errorf("validateClientRegistrationURLs(): got nil error, want error for dangerous %s", tt.name) + } else if !strings.Contains(err.Error(), tt.name) { + t.Errorf("validateClientRegistrationURLs(): got error %v, want error containing %q", err, tt.name) + } + }) + } + + t.Run("empty optional fields are valid", func(t *testing.T) { + meta := &ClientRegistrationMetadata{ + RedirectURIs: []string{"https://app.example.com/callback"}, + // All optional URL fields left empty + } + if err := validateClientRegistrationURLs(meta); err != nil { + t.Errorf("validateClientRegistrationURLs(): got err %v, want nil", err) + } + }) +} + +// TestGetAuthServerMetaRejectsDangerousURLs tests that GetAuthServerMeta rejects +// metadata containing dangerous URL schemes. +func TestGetAuthServerMetaRejectsDangerousURLs(t *testing.T) { + tests := []struct { + name string + metadata AuthServerMeta + wantErrText string + }{ + { + name: "javascript authorization_endpoint", + metadata: AuthServerMeta{ + Issuer: "", // Will be set dynamically + AuthorizationEndpoint: "javascript:alert('xss')", + TokenEndpoint: "https://auth.example.com/token", + JWKSURI: "https://auth.example.com/.well-known/jwks.json", + ResponseTypesSupported: []string{"code"}, + CodeChallengeMethodsSupported: []string{"S256"}, + }, + wantErrText: "authorization_endpoint", + }, + { + name: "data token_endpoint", + metadata: AuthServerMeta{ + Issuer: "", + AuthorizationEndpoint: "https://auth.example.com/authorize", + TokenEndpoint: "data:text/html,", + JWKSURI: "https://auth.example.com/.well-known/jwks.json", + ResponseTypesSupported: []string{"code"}, + CodeChallengeMethodsSupported: []string{"S256"}, + }, + wantErrText: "token_endpoint", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + meta := tt.metadata + meta.Issuer = "https://" + r.Host + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(meta) + })) + defer server.Close() + + ctx := context.Background() + issuer := server.URL + metadataURL := issuer + _, err := GetAuthServerMeta(ctx, metadataURL, issuer, server.Client()) + if err == nil { + t.Fatal("GetAuthServerMeta(): got nil error, want error") + } + if !strings.Contains(err.Error(), tt.wantErrText) { + t.Errorf("GetAuthServerMeta(): got error %v, want error containing %q", err, tt.wantErrText) + } + }) + } +} + +// TestGetProtectedResourceMetadataRejectsDangerousURLs tests that +// GetProtectedResourceMetadata rejects metadata with dangerous authorization server URLs. +func TestGetProtectedResourceMetadataRejectsDangerousURLs(t *testing.T) { + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + serverURL := "https://" + r.Host + meta := ProtectedResourceMetadata{ + Resource: serverURL, + AuthorizationServers: []string{"javascript:alert('xss')"}, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(meta) + })) + defer server.Close() + + ctx := context.Background() + metadataURL := server.URL + "/.well-known/oauth-protected-resource" + _, err := GetProtectedResourceMetadata(ctx, metadataURL, server.URL, server.Client()) + if err == nil { + t.Fatal("GetProtectedResourceMetadata(): got nil error, want error") + } + if !strings.Contains(err.Error(), "disallowed scheme") { + t.Errorf("GetProtectedResourceMetadata(): got error %v, want error containing \"disallowed scheme\"", err) + } +} + +// TestRegisterClientRejectsDangerousURLs tests that RegisterClient rejects +// responses containing dangerous URL schemes. +func TestRegisterClientRejectsDangerousURLs(t *testing.T) { + tests := []struct { + name string + responseJSON string + wantErrText string + }{ + { + name: "javascript redirect_uri in response", + responseJSON: `{ + "client_id": "test-client", + "redirect_uris": ["javascript:alert(1)"] + }`, + wantErrText: "redirect_uris[0]", + }, + { + name: "data client_uri", + responseJSON: `{ + "client_id": "test-client", + "redirect_uris": ["https://app.example.com/callback"], + "client_uri": "data:text/html," + }`, + wantErrText: "client_uri", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusCreated) + w.Write([]byte(tt.responseJSON)) + })) + defer server.Close() + + ctx := context.Background() + clientMeta := &ClientRegistrationMetadata{ + RedirectURIs: []string{"https://app.example.com/callback"}, + } + + _, err := RegisterClient(ctx, server.URL+"/register", clientMeta, server.Client()) + if err == nil { + t.Fatal("RegisterClient(): got nil error, want error") + } + if !strings.Contains(err.Error(), tt.wantErrText) { + t.Errorf("RegisterClient(): got error %v, want error containing %q", err, tt.wantErrText) + } + }) + } +} diff --git a/scripts/client-conformance.sh b/scripts/client-conformance.sh new file mode 100755 index 0000000..fef74ca --- /dev/null +++ b/scripts/client-conformance.sh @@ -0,0 +1,92 @@ +#!/bin/bash +# Copyright 2025 The Go MCP SDK Authors. All rights reserved. +# Use of this source code is governed by an MIT-style +# license that can be found in the LICENSE file. + +# Run MCP conformance tests against the Go SDK conformance client. + +set -e + +RESULT_DIR="" +WORKDIR="" +CONFORMANCE_REPO="" +SUITE="core" +FINAL_EXIT_CODE=0 + +usage() { + echo "Usage: $0 [options]" + echo "" + echo "Run MCP conformance tests against the Go SDK conformance client." + echo "" + echo "Options:" + echo " --result_dir Save results to the specified directory" + echo " --conformance_repo Run conformance tests from a local checkout" + echo " instead of using the latest npm release" + echo " --suite Which suite to run (default: core)" + echo " --help Show this help message" +} + + +# Parse arguments. +while [[ $# -gt 0 ]]; do + case $1 in + --result_dir) + RESULT_DIR="$2" + shift 2 + ;; + --conformance_repo) + CONFORMANCE_REPO="$2" + shift 2 + ;; + --suite) + SUITE="$2" + shift 2 + ;; + --help) + usage + exit 0 + ;; + *) + echo "Unknown option: $1" + usage + exit 1 + ;; + esac +done + +# Set up the work directory. +if [ -n "$RESULT_DIR" ]; then + mkdir -p "$RESULT_DIR" + WORKDIR="$RESULT_DIR" +else + WORKDIR=$(mktemp -d) +fi + +# Build the conformance server. +go build -o "$WORKDIR/conformance-client" ./conformance/everything-client + +# Run conformance tests from the work directory to avoid writing results to the repo. +echo "Running conformance tests..." +if [ -n "$CONFORMANCE_REPO" ]; then + # Run from local checkout using npm run start. + (cd "$WORKDIR" && \ + npm --prefix "$CONFORMANCE_REPO" run start -- \ + client --command "$WORKDIR/conformance-client" \ + --suite "$SUITE" \ + ${RESULT_DIR:+--output-dir "$RESULT_DIR"}) || FINAL_EXIT_CODE=$? +else + (cd "$WORKDIR" && \ + npx @modelcontextprotocol/conformance@latest \ + client --command "$WORKDIR/conformance-client" \ + --suite "$SUITE" \ + ${RESULT_DIR:+--output-dir "$RESULT_DIR"}) || FINAL_EXIT_CODE=$? +fi + +echo "" +if [ -n "$RESULT_DIR" ]; then + echo "See $RESULT_DIR for details." +else + echo "Run with --result_dir to save results." +fi + +exit $FINAL_EXIT_CODE diff --git a/scripts/server-conformance.sh b/scripts/server-conformance.sh new file mode 100755 index 0000000..e8fe6e1 --- /dev/null +++ b/scripts/server-conformance.sh @@ -0,0 +1,115 @@ +#!/bin/bash +# Copyright 2025 The Go MCP SDK Authors. All rights reserved. +# Use of this source code is governed by an MIT-style +# license that can be found in the LICENSE file. + +# Run MCP conformance tests against the Go SDK conformance server. + +set -e + +PORT="${PORT:-3000}" +SERVER_PID="" +RESULT_DIR="" +WORKDIR="" +CONFORMANCE_REPO="" +FINAL_EXIT_CODE=0 + +usage() { + echo "Usage: $0 [options]" + echo "" + echo "Run MCP conformance tests against the Go SDK conformance server." + echo "" + echo "Options:" + echo " --result_dir Save results to the specified directory" + echo " --conformance_repo Run conformance tests from a local checkout" + echo " instead of using the latest npm release" + echo " --help Show this help message" +} + +# Parse arguments. +while [[ $# -gt 0 ]]; do + case $1 in + --result_dir) + RESULT_DIR="$2" + shift 2 + ;; + --conformance_repo) + CONFORMANCE_REPO="$2" + shift 2 + ;; + --help) + usage + exit 0 + ;; + *) + echo "Unknown option: $1" + usage + exit 1 + ;; + esac +done + +cleanup() { + if [ -n "$SERVER_PID" ]; then + echo "Stopping server..." + kill "$SERVER_PID" 2>/dev/null || true + wait "$SERVER_PID" 2>/dev/null || true + fi +} +trap cleanup EXIT + +# Set up the work directory. +if [ -n "$RESULT_DIR" ]; then + mkdir -p "$RESULT_DIR" + WORKDIR="$RESULT_DIR" +else + WORKDIR=$(mktemp -d) +fi + +# Build the conformance server. +go build -o "$WORKDIR/conformance-server" ./conformance/everything-server + +# Start the server in the background. +# -stateless=false pins the server to the stateful transport so that +# server-initiated sampling/elicitation scenarios (which are not supported on +# stateless streamable HTTP) work against the current @latest conformance +# suite. Drop this flag once the 0.2.x line is promoted to the @latest +# dist-tag on npm and the stateless leg becomes viable. +echo "Starting conformance server on localhost:$PORT..." +"$WORKDIR/conformance-server" -http="localhost:$PORT" -stateless=false & +SERVER_PID=$! + +echo "Server pid is $SERVER_PID" + +# Wait for server to be ready +echo "Waiting for server to be ready..." +if ! timeout 15 bash -c "until curl -s http://localhost:$PORT > /dev/null 2>&1; do sleep 0.5; done"; then + echo "Server failed to start within 15 seconds." + exit 1 +fi + +# Run conformance tests from the work directory to avoid writing results to the repo. +echo "Running conformance tests..." +if [ -n "$CONFORMANCE_REPO" ]; then + # Run from local checkout using npm run start. + (cd "$WORKDIR" && \ + npm --prefix "$CONFORMANCE_REPO" run start -- \ + server --url "http://localhost:$PORT" \ + --spec-version 2025-11-25 \ + ${RESULT_DIR:+--output-dir "$RESULT_DIR"}) || FINAL_EXIT_CODE=$? +else + (cd "$WORKDIR" && \ + npx @modelcontextprotocol/conformance@latest \ + server --url "http://localhost:$PORT" \ + --spec-version 2025-11-25 \ + ${RESULT_DIR:+--output-dir "$RESULT_DIR"}) || FINAL_EXIT_CODE=$? +fi + +echo "" +if [ -n "$RESULT_DIR" ]; then + echo "See $RESULT_DIR for details." +else + echo "Run with --result_dir to save results." +fi + +exit $FINAL_EXIT_CODE